Spaces:
Running
Running
File size: 7,855 Bytes
09ca579 b8bfb45 09ca579 b8bfb45 09ca579 469696e b8bfb45 469696e b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 469696e 09ca579 1eacb01 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 469696e 1eacb01 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 b8bfb45 09ca579 469696e b8bfb45 469696e 09ca579 b8bfb45 09ca579 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 |
#!/usr/bin/env python
from __future__ import annotations
import pathlib
import shlex
import subprocess
import tempfile
import gradio as gr
from omegaconf import OmegaConf
def gen_feature_extraction_config(
exp_name: str,
prompt: str,
seed: int,
guidance_scale: float,
ddim_steps: int,
) -> str:
config = OmegaConf.load("plug-and-play/configs/pnp/feature-extraction-generated.yaml")
config.config.experiment_name = exp_name
config.config.prompt = prompt
config.config.seed = seed
config.config.scale = guidance_scale
config.config.ddim_steps = ddim_steps
temp_file = tempfile.NamedTemporaryFile(suffix=".yaml", delete=False)
with open(temp_file.name, "w") as f:
f.write(OmegaConf.to_yaml(config))
return temp_file.name
def run_feature_extraction_command(
prompt: str,
seed: int,
guidance_scale: float,
ddim_steps: int,
) -> tuple[str, str]:
exp_name = f'{prompt.replace(" ", "_")}_{seed}_{guidance_scale:.1f}_{ddim_steps}'
if not pathlib.Path(f"plug-and-play/experiments/{exp_name}").exists():
config_path = gen_feature_extraction_config(
exp_name,
prompt,
seed,
guidance_scale,
ddim_steps,
)
subprocess.run(shlex.split(f"python run_features_extraction.py --config {config_path}"), cwd="plug-and-play")
return f"plug-and-play/experiments/{exp_name}/samples/0.png", exp_name
def gen_pnp_config(
exp_name: str,
prompt: str,
guidance_scale: float,
ddim_steps: int,
feature_injection_threshold: int,
negative_prompt: str,
negative_prompt_alpha: float,
negative_prompt_schedule: str,
) -> str:
config = OmegaConf.load("plug-and-play/configs/pnp/pnp-generated.yaml")
config.source_experiment_name = exp_name
config.prompts = [prompt]
config.scale = guidance_scale
config.num_ddim_sampling_steps = ddim_steps
config.feature_injection_threshold = feature_injection_threshold
config.negative_prompt = negative_prompt
config.negative_prompt_alpha = negative_prompt_alpha
config.negative_prompt_schedule = negative_prompt_schedule
temp_file = tempfile.NamedTemporaryFile(suffix=".yaml", delete=False)
with open(temp_file.name, "w") as f:
f.write(OmegaConf.to_yaml(config))
return temp_file.name
def run_pnp_command(
exp_name: str,
prompt: str,
negative_prompt: str,
guidance_scale: float,
ddim_steps: int,
feature_injection_threshold: int,
negative_prompt_alpha: float,
negative_prompt_schedule: str,
) -> str:
config_path = gen_pnp_config(
exp_name,
prompt,
guidance_scale,
ddim_steps,
feature_injection_threshold,
negative_prompt,
negative_prompt_alpha,
negative_prompt_schedule,
)
subprocess.run(shlex.split(f"python run_pnp.py --config {config_path}"), cwd="plug-and-play")
out_dir = pathlib.Path(
f'plug-and-play/experiments/{exp_name}/translations/{guidance_scale}_{prompt.replace(" ", "_")}'
)
out_label = f'INJECTION_T_{feature_injection_threshold}_STEPS_{ddim_steps}_NP-ALPHA_{negative_prompt_alpha}_SCHEDULE_{negative_prompt_schedule}_NP_{negative_prompt.replace(" ", "_")}'
out_path = out_dir / f"{out_label}_sample_0.png"
return out_path.as_posix()
def process_example(source_prompt: str, seed: int, translation_prompt: str) -> tuple[str, str, str]:
generated_image, exp_name = run_feature_extraction_command(source_prompt, seed, guidance_scale=5, ddim_steps=50)
result = run_pnp_command(
exp_name,
translation_prompt,
negative_prompt="",
guidance_scale=7.5,
ddim_steps=50,
feature_injection_threshold=40,
negative_prompt_alpha=0.75,
negative_prompt_schedule="linear",
)
return generated_image, exp_name, result
def create_prompt_demo() -> gr.Blocks:
with gr.Blocks() as demo:
with gr.Group():
gr.Markdown("Step 1 (This step will take about 1.5 minutes on A10G.)")
with gr.Row():
with gr.Column():
source_prompt = gr.Text(label="Source prompt")
seed = gr.Slider(label="Seed", minimum=0, maximum=100000, step=1, value=0)
with gr.Accordion(label="Advanced settings", open=False):
source_guidance_scale = gr.Slider(
label="Guidance scale", minimum=0, maximum=50, step=0.1, value=5
)
source_ddim_steps = gr.Slider(label="DDIM steps", minimum=1, maximum=100, step=1, value=50)
extract_feature_button = gr.Button("Generate and extract features")
with gr.Column():
generated_image = gr.Image(label="Generated image", type="filepath")
exp_name = gr.Text(visible=False)
with gr.Group():
gr.Markdown("Step 2 (This step will take about 1.5 minutes on A10G.)")
with gr.Row():
with gr.Column():
translation_prompt = gr.Text(label="Prompt for translation")
negative_prompt = gr.Text(label="Negative prompt")
with gr.Accordion(label="Advanced settings", open=False):
guidance_scale = gr.Slider(label="Guidance scale", minimum=0, maximum=50, step=0.1, value=7.5)
ddim_steps = gr.Slider(
label="Number of inference steps", minimum=1, maximum=100, step=1, value=50
)
feature_injection_threshold = gr.Slider(
label="Feature injection threshold", minimum=0, maximum=100, step=1, value=40
)
negative_prompt_alpha = gr.Slider(
label="Negative prompt alpha", minimum=0, maximum=1, step=0.01, value=0.75
)
negative_prompt_schedule = gr.Dropdown(
label="Negative prompt schedule", choices=["linear", "constant", "exp"], value="linear"
)
generate_button = gr.Button("Generate")
with gr.Column():
result = gr.Image(label="Result", type="filepath")
with gr.Row():
gr.Examples(
examples=[
["horse in mud", 50, "a photo of a zebra in the snow"],
["horse in mud", 50, "a photo of a husky in the grass"],
],
inputs=[
source_prompt,
seed,
translation_prompt,
],
outputs=[
generated_image,
exp_name,
result,
],
fn=process_example,
)
extract_feature_button.click(
fn=run_feature_extraction_command,
inputs=[
source_prompt,
seed,
source_guidance_scale,
source_ddim_steps,
],
outputs=[
generated_image,
exp_name,
],
)
generate_button.click(
fn=run_pnp_command,
inputs=[
exp_name,
translation_prompt,
negative_prompt,
guidance_scale,
ddim_steps,
feature_injection_threshold,
negative_prompt_alpha,
negative_prompt_schedule,
],
outputs=result,
)
return demo
if __name__ == "__main__":
demo = create_prompt_demo()
demo.queue().launch()
|