Advanced settings: pick v0.2.1 instead of v0.3 (both LoRAs loaded, switched per call)

app.py

@@ -1,4 +1,5 @@
-"""Gradio demo for the distilled Qwen-Image-2.1 turbo (v0.3: LoRA r256 sampled in 6 steps; 9 steps = 7 turbo + 2 base-model steps)."""
+"""Gradio demo for the distilled Qwen-Image-2.1 turbo (v0.3: LoRA r256 sampled in 6 steps; 9 steps = 7 turbo + 2 base-model steps;
+v0.2.1 selectable under Advanced settings)."""
 
 # ZeroGPU patches torch at import time, so `spaces` must be imported before torch or any CUDA use.
 # find_spec keeps this file runnable off-Spaces, where the package is absent.
@@ -38,13 +39,17 @@ from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline
 from diffusers.pipelines.qwenimage21.pipeline_qwenimage21 import calculate_dimensions
 from huggingface_hub import CommitScheduler
 
-# The v0.3 LoRA (r256) at the model repo root, applied at runtime on top of the base transformer and never merged
+# The LoRAs (r256) at the model repo root, applied at runtime on top of the base transformer and never merged
 # (merging into bf16 keeps only ~47% of the adapter delta, the per-element deltas sit below the bf16 ULP of the base
-# weights). LORA_FILE selects the adapter file; the v0.2.1 and v0.2 adapters are still in the repo.
+# weights). v0.3 (the default) and v0.2.1 are both loaded as named adapters (PEFT adapter names cannot contain dots)
+# and each call activates the one picked under Advanced settings; both are sampled on the same nodes (raw_nodes below).
 BASE_MODEL_ID = os.environ.get("BASE_MODEL_ID", "Qwen/Qwen-Image-2.1")
 STUDENT_REPO = os.environ.get("STUDENT_REPO", os.environ.get("LORA_REPO", "Viggle/Qwen-Image-2.1-viggle-turbo"))
-LORA_FILE = os.environ.get("LORA_FILE", "Qwen-Image-2.1-viggle-turbo-v0.3-6step-lora-r256.safetensors")
-STUDENT_TAG = "v0.3 LoRA r256" if "v0.3" in LORA_FILE else f"LoRA {LORA_FILE}"
+LORAS = {  # version -> (adapter name, file)
+    "v0.3": ("v03", "Qwen-Image-2.1-viggle-turbo-v0.3-6step-lora-r256.safetensors"),
+    "v0.2.1": ("v021", "Qwen-Image-2.1-viggle-turbo-v0.2.1-6step-lora-r256.safetensors"),
+}
+VERSION = "v0.3"
 HF_TOKEN = os.environ.get("HF_TOKEN")  # only needed while the model repo is private
 # Usage counter. The Space autoscales to several replicas and its log view shows only one of them, so the [gen] log lines
 # undercount. Every generate() call that reaches the GPU also appends one JSON line (metadata only: no prompt, no images)
@@ -134,7 +139,8 @@ CUSTOM_CAP = {False: 2048, True: 1536}  # keyed by "has references"
 MAX_SIDE = max(max(size) for size in SIZES.values() if size)
 
 pipe = QwenImage21Pipeline.from_pretrained(BASE_MODEL_ID, dtype=torch.bfloat16, token=HF_TOKEN)
-pipe.load_lora_weights(STUDENT_REPO, weight_name=LORA_FILE, token=HF_TOKEN)
+for adapter, lora_file in LORAS.values():
+    pipe.load_lora_weights(STUDENT_REPO, weight_name=lora_file, adapter_name=adapter, token=HF_TOKEN)
 # The stock Qwen-Image-2.1 scheduler config carries shift_terminal=0.02, which stretches the last
 # sigma node to 0.02 instead of 0 and silently costs the final step. The student was distilled
 # against the unstretched schedule, so rebuild the scheduler with shift_terminal=None.
@@ -233,11 +239,12 @@ def base_tail(pipe, i, t, kwargs):
 
 
 @gpu
-def denoise(prompt, images, width, height, steps, seed):
+def denoise(prompt, images, width, height, steps, seed, version=VERSION):
     # No VAE tiling anywhere (xlarge card): tiling the reference encode wrecks edits, and tiled decodes are not the
     # validated pipeline either.
     generator = torch.Generator(device="cuda").manual_seed(seed)
     started = time.perf_counter()
+    pipe.set_adapters(LORAS[version][0])
     pipe.enable_lora()  # every call starts on the student, even if an earlier 9-step call died with the LoRA off
     reextract[0] = False
     result = pipe(
@@ -256,8 +263,10 @@ def denoise(prompt, images, width, height, steps, seed):
     return result, time.perf_counter() - started
 
 
-def generate(prompt, refs, size_label, seed, randomize_seed, enhance="Auto", steps=STEPS, custom_width=1024, custom_height=1024):
+def generate(prompt, refs, size_label, seed, randomize_seed, enhance="Auto", steps=STEPS, custom_width=1024, custom_height=1024,
+             version=VERSION):
     t0 = time.perf_counter()
+    version = version if version in LORAS else VERSION
     steps = min(max(int(steps), 3), 9)  # 6 (RAW_NODES) is the default; 9 = 7 student steps + 2 base-model steps; fewer than 6 degrade quickly
     # refs: the gallery's (filepath, caption) pairs in upload order, which is the order the prompt's "image 1, 2, ..." refer
     # to. EXIF rotation + RGB is what the gr.Image slots used to do.
@@ -294,18 +303,18 @@ def generate(prompt, refs, size_label, seed, randomize_seed, enhance="Auto", ste
         enhance_note = f" · enhanced {pe_time:.1f}s" if used_prompt != prompt else " · enhancement failed, prompt sent as written"
     mode = "edit" if images else "t2i"
     try:
-        result, elapsed = denoise(used_prompt, images, width, height, steps, seed)
+        result, elapsed = denoise(used_prompt, images, width, height, steps, seed, version)
     except Exception as e:  # ZeroGPU quota / GPU timeout errors land here: count them, then show them as before
-        record(ok=False, mode=mode, refs=len(images), steps=steps, error=f"{type(e).__name__}: {e}"[:200],
+        record(ok=False, mode=mode, refs=len(images), steps=steps, version=version, error=f"{type(e).__name__}: {e}"[:200],
                total=round(time.perf_counter() - t0, 2))
         raise
     # one line per generation, for counting usage from the Space logs (no prompt text)
     print(f"[gen] {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} mode={mode} refs={len(images)} "
-          f"steps={steps} size={result.width}x{result.height} enhance={int(used_prompt != prompt)} pe={pe_time:.1f}s gen={elapsed:.2f}s",
+          f"steps={steps} version={version} size={result.width}x{result.height} enhance={int(used_prompt != prompt)} pe={pe_time:.1f}s gen={elapsed:.2f}s",
           flush=True)
-    record(ok=True, mode=mode, refs=len(images), steps=steps, size=f"{result.width}x{result.height}", enhance=enhance,
+    record(ok=True, mode=mode, refs=len(images), steps=steps, version=version, size=f"{result.width}x{result.height}", enhance=enhance,
            enhanced=used_prompt != prompt, pe=round(pe_time, 2), gen=round(elapsed, 2), total=round(time.perf_counter() - t0, 2))
-    return result, f"seed `{seed}` · {result.width}×{result.height} · {steps} steps · {elapsed:.2f}s · {STUDENT_TAG}{enhance_note}{clamped}", used_prompt
+    return result, f"seed `{seed}` · {result.width}×{result.height} · {steps} steps · {elapsed:.2f}s · {version} LoRA r256{enhance_note}{clamped}", used_prompt
 
 
 def refresh_sizes(prompt, refs, current):
@@ -397,8 +406,8 @@ with gr.Blocks(title="Viggle Turbo v0.3 · Qwen-Image-2.1 6-step", theme=THEME,
             "**v0.3 (2026-09-29):** at 6 steps, less grain than v0.2.1 and a little softer on fine texture. We think 6 steps "
             "is close to its capacity: every further gain we found cost something elsewhere. **9 steps** runs 7 turbo "
             "steps and lets the base model finish the last two: finer detail and small text right more often, at about 1.4–1.5× the "
-            "time of 6 steps. "
-            f"Weights: **{STUDENT_TAG}** from [{STUDENT_REPO}](https://huggingface.co/{STUDENT_REPO})."
+            "time of 6 steps. v0.2.1 can still be picked under **Advanced settings**. "
+            f"Weights: **v0.3 LoRA r256** (and v0.2.1) from [{STUDENT_REPO}](https://huggingface.co/{STUDENT_REPO})."
         )
     with gr.Tabs() as tabs:
         with gr.Tab("Generate", id="generate"):
@@ -432,6 +441,9 @@ with gr.Blocks(title="Viggle Turbo v0.3 · Qwen-Image-2.1 6-step", theme=THEME,
                     with gr.Accordion("Advanced settings", open=False) as advanced:
                         steps = gr.Slider(label="Steps", minimum=3, maximum=9, step=1, value=STEPS,
                                           info="6 is the default; 7-8 add steps at the high-noise end; 9 = 7 turbo + 2 base-model steps (finer detail; ~1.4–1.5× the time).")
+                        version = gr.Radio(list(LORAS), value=VERSION, label="Model version",
+                                           info="v0.3 is the default. v0.2.1 is the previous release: a little sharper on fine "
+                                           "texture, with more grain.")
                         with gr.Row():
                             seed = gr.Number(label="Seed", value=0, precision=0, interactive=False)
                             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
@@ -446,12 +458,13 @@ with gr.Blocks(title="Viggle Turbo v0.3 · Qwen-Image-2.1 6-step", theme=THEME,
             def example_details(prompt, refs, result):
                 """Runs on an example click (no GPU): fills the status line and the prompt stored with the row, and sets the
                 steps / size / enhancement / seed controls to what produced it, so Generate reproduces the result. It opens
-                Advanced settings to show that the seed is now fixed. It also sets the size menu's mode itself: a
-                text-to-image row loaded after another one leaves the gallery unchanged, so refresh_sizes does not run."""
+                Advanced settings to show that the seed is now fixed, and picks v0.3, which rendered every example. It also sets
+                the size menu's mode itself: a text-to-image row loaded after another one leaves the gallery unchanged, so
+                refresh_sizes does not run."""
                 row = next(row for row in EXAMPLES if row["prompt"] == prompt)
                 choices = EDIT_CHOICES if any(row["refs"]) else T2I_CHOICES
                 return (row["info"], row["used_prompt"], row["steps"], gr.update(choices=choices, value=row["size_label"]),
-                        "On" if row["enhance"] else "Off", row["seed"], False, gr.update(open=True))
+                        "On" if row["enhance"] else "Off", row["seed"], False, VERSION, gr.update(open=True))
 
             if EXAMPLES:
                 gr.Examples(
@@ -461,7 +474,7 @@ with gr.Blocks(title="Viggle Turbo v0.3 · Qwen-Image-2.1 6-step", theme=THEME,
                     ],
                     inputs=[prompt, refs, output_image],
                     fn=example_details,
-                    outputs=[info, used_prompt, steps, size_label, enhance, seed, randomize_seed, advanced],
+                    outputs=[info, used_prompt, steps, size_label, enhance, seed, randomize_seed, version, advanced],
                     run_on_click=True,
                     # On Spaces gr.Examples caches by default (lazily on ZeroGPU), and gradio 5.50 crashes caching an output that
                     # updates a Dropdown's choices ('Dropdown' object has no attribute 'proxy_url'). example_details needs no GPU.
@@ -491,7 +504,7 @@ with gr.Blocks(title="Viggle Turbo v0.3 · Qwen-Image-2.1 6-step", theme=THEME,
     randomize_seed.change(lambda on: gr.update(interactive=not on), inputs=randomize_seed, outputs=seed)
     run.click(
         generate,
-        inputs=[prompt, refs, size_label, seed, randomize_seed, enhance, steps, custom_width, custom_height],
+        inputs=[prompt, refs, size_label, seed, randomize_seed, enhance, steps, custom_width, custom_height, version],
         outputs=[output_image, info, used_prompt],
     )
     demo.load(open_tab, outputs=tabs, show_progress="hidden")