Add guidance schedule controls

app.py

@@ -36,10 +36,12 @@ MAX_SEED = 2**31 - 1
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 MODES = {
-    "Turbo · 12 steps": dict(num_inference_steps=12, guidance_schedule=(7.0,) * 11 + (3.0,) * 1, mu=0.5, std=1.75),
-    "Default · 20 steps": dict(num_inference_steps=20, guidance_schedule=(7.0,) * 18 + (3.0,) * 2, mu=0.0, std=1.75),
-    "Quality · 48 steps": dict(num_inference_steps=48, guidance_schedule=(7.0,) * 45 + (3.0,) * 3, mu=0.0, std=1.5),
+    "Turbo · 12 steps": dict(num_inference_steps=12, final_guidance_steps=1, mu=0.5, std=1.75),
+    "Default · 20 steps": dict(num_inference_steps=20, final_guidance_steps=2, mu=0.0, std=1.75),
+    "Quality · 48 steps": dict(num_inference_steps=48, final_guidance_steps=3, mu=0.0, std=1.5),
 }
+DEFAULT_MAIN_GUIDANCE = 7.0
+DEFAULT_FINAL_GUIDANCE = 3.0
 
 DEFAULT_CAPTION = {
     "high_level_description": "A clean poster announcing a small experimental image generation lab.",
@@ -97,6 +99,16 @@ def normalize_caption(raw_caption):
     return json.dumps(caption, ensure_ascii=False, separators=(",", ":")), caption
 
 
+def build_preset(mode, main_guidance=DEFAULT_MAIN_GUIDANCE, final_guidance=DEFAULT_FINAL_GUIDANCE):
+    preset = dict(MODES.get(mode, MODES["Default · 20 steps"]))
+    steps = int(preset.pop("num_inference_steps"))
+    final_steps = min(int(preset.pop("final_guidance_steps")), steps)
+    main_steps = steps - final_steps
+    guidance_schedule = (float(main_guidance),) * main_steps + (float(final_guidance),) * final_steps
+    preset.update(num_inference_steps=steps, guidance_schedule=guidance_schedule)
+    return preset
+
+
 t = time.perf_counter()
 if TEXT_ENCODER_ID:
     text_encoder = AutoModel.from_pretrained(
@@ -161,32 +173,46 @@ def _per_step(width, height):
     return max(0.2, _PS_A + _PS_B * ((int(width) // 16) * (int(height) // 16)))
 
 
-def _gpu_duration(caption_text, mode, width, height, seed, progress=None):
+def _gpu_duration(caption_text, mode, width, height, seed, main_guidance, final_guidance, progress=None):
     steps = MODES.get(mode, MODES["Default · 20 steps"])["num_inference_steps"]
     budget = steps * _per_step(width, height) + DIFFUSION_OVERHEAD_S
     return max(60, int(math.ceil(budget * DURATION_MARGIN)))
 
 
 @spaces.GPU(duration=_gpu_duration, size="xlarge")
-def _gpu_generate(caption_text, mode, width, height, seed, progress=gr.Progress(track_tqdm=True)):
+def _gpu_generate(caption_text, mode, width, height, seed, main_guidance, final_guidance, progress=gr.Progress(track_tqdm=True)):
     aoti_thread = Thread(target=_apply_aoti, daemon=True)
     aoti_thread.start()
     aoti_thread.join()
 
     progress(0.0, desc="Generating image")
     generator = torch.Generator(device="cuda").manual_seed(int(seed))
-    preset = MODES.get(mode, MODES["Default · 20 steps"])
+    preset = build_preset(mode, main_guidance, final_guidance)
     t = time.perf_counter()
     image = pipe(prompt=caption_text, width=int(width), height=int(height), generator=generator, **preset).images[0]
-    print(f"[timing] diffusion ({mode}): {time.perf_counter() - t:.2f}s", flush=True)
+    print(
+        f"[timing] diffusion ({mode}, guidance={float(main_guidance):.2f}->{float(final_guidance):.2f}): "
+        f"{time.perf_counter() - t:.2f}s",
+        flush=True,
+    )
     return image
 
 
-def generate(caption_json, mode="Default · 20 steps", width=1024, height=1024, seed=0, randomize_seed=False, progress=gr.Progress(track_tqdm=True)):
+def generate(
+    caption_json,
+    mode="Default · 20 steps",
+    width=1024,
+    height=1024,
+    seed=0,
+    randomize_seed=False,
+    main_guidance=DEFAULT_MAIN_GUIDANCE,
+    final_guidance=DEFAULT_FINAL_GUIDANCE,
+    progress=gr.Progress(track_tqdm=True),
+):
     caption_text, parsed_caption = normalize_caption(caption_json)
     if randomize_seed or seed < 0:
         seed = random.randint(0, MAX_SEED)
-    image = _gpu_generate(caption_text, mode, width, height, seed)
+    image = _gpu_generate(caption_text, mode, width, height, seed, main_guidance, final_guidance)
     return image, int(seed), parsed_caption, caption_text
 
 
@@ -209,6 +235,9 @@ with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4 JSON Lab", css=CSS) a
                 mode = gr.Radio(choices=list(MODES.keys()), value="Default · 20 steps", label="Mode")
                 width = gr.Slider(512, 2048, value=1024, step=64, label="Width")
                 height = gr.Slider(512, 2048, value=1024, step=64, label="Height")
+            with gr.Row():
+                main_guidance = gr.Slider(0.0, 9.0, value=DEFAULT_MAIN_GUIDANCE, step=0.25, label="Main guidance")
+                final_guidance = gr.Slider(0.0, 9.0, value=DEFAULT_FINAL_GUIDANCE, step=0.25, label="Final guidance")
             with gr.Row():
                 seed = gr.Number(label="Seed", value=0, precision=0)
                 randomize = gr.Checkbox(label="Randomize seed", value=False)
@@ -266,7 +295,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4 JSON Lab", css=CSS) a
 
     run.click(
         generate,
-        inputs=[caption, mode, width, height, seed, randomize],
+        inputs=[caption, mode, width, height, seed, randomize, main_guidance, final_guidance],
         outputs=[out_image, seed, out_caption, out_text],
     )