feat: update with new safeternsors

README.md

@@ -1,14 +1,37 @@
 ---
-title: Pruna Qwen Image 2.1
+title: P-Q-Image-2.1
 emoji: ⚡
 colorFrom: indigo
 colorTo: indigo
 sdk: gradio
 sdk_version: 6.28.0
+python_version: "3.12"
 app_file: app.py
 pinned: false
 hf_oauth: true
 hf_oauth_scopes:
 - inference-api
 - write-repos
+short_description: Few-step LoRA for Qwen-Image-2.1 (6-steps v0.2)
+models:
+  - PrunaAI/P-Q-Image-2.1
+  - Qwen/Qwen-Image-2.1
+license: other
+license_name: qwen-research
 ---
+
+# P-Q-Image-2.1
+
+Few-step LoRA adapters for [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1).
+Same pipeline for text-to-image and image editing, no CFG.
+
+Pick a **Variant**:
+
+| Variant | Steps | Notes |
+|---|:---:|---|
+| `6-steps` | 6 | v0.2 default. Faithful few-step distillation; best for editing. |
+| `6-steps-upsampled` | 6 | v0.2. Distills prompt upsampling into the LoRA. |
+| `8-steps` | 8 | v0.1, higher quality of that release. |
+| `5-steps` | 5 | v0.1, faster, lower quality. |
+
+Model card: [`PrunaAI/P-Q-Image-2.1`](https://huggingface.co/PrunaAI/P-Q-Image-2.1)

app.py

@@ -1,52 +1,139 @@
-"""gr.Workflow app for PrunaAI/Pruna-Qwen-Image-2.1 (ZeroGPU).
+"""gr.Workflow app for PrunaAI/P-Q-Image-2.1 (ZeroGPU).
 
 Nodes on the canvas:
-  - "Generate Image": text-to-image with the 5- or 8-step LoRA adapter
+  - "Generate Image": text-to-image with a v0.2 or v0.1 LoRA adapter
   - "Edit Image": image editing with an input image + prompt
 """
 
 import os
-
-import torch
 import spaces
+import torch
 import gradio as gr
 from gradio.utils import get_upload_folder
 from gradio_client import utils as client_utils
 from PIL import Image
 from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline
 
-SIGMAS = {
-    5: [1.0, 0.94, 6 / 7, 2 / 3, 0.4],
-    8: [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
+LORA_REPO = "PrunaAI/P-Q-Image-2.1"
+DEFAULT_VARIANT = "6-steps"
+LORA_STRENGTH = 1.0
+
+# Each adapter is trained for one schedule. v0.2 uses unshifted 6-step sigmas
+# plus dynamic shifting; v0.1 uses its raw nodes with shifting off.
+VARIANTS = {
+    "6-steps": {
+        "weight_name": "p_q_image_2.1_6_steps_v0.2.safetensors",
+        "adapter_name": "step6_v02",
+        "steps": 6,
+        "sigmas": [1.0, 0.925, 0.85, 0.625, 0.375, 0.125],
+        "dynamic_shifting": True,
+        "use_kv_cache": False,
+    },
+    "6-steps-upsampled": {
+        "weight_name": "p_q_image_2.1_6_steps_upsampled_v0.2.safetensors",
+        "adapter_name": "step6_upsampled_v02",
+        "steps": 6,
+        "sigmas": [1.0, 0.925, 0.85, 0.625, 0.375, 0.125],
+        "dynamic_shifting": True,
+        "use_kv_cache": False,
+    },
+    "8-steps": {
+        "weight_name": "p_q_image_2.1_8_steps_v0.1.safetensors",
+        "adapter_name": "step8_v01",
+        "steps": 8,
+        "sigmas": [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
+        "dynamic_shifting": False,
+        "use_kv_cache": True,
+    },
+    "5-steps": {
+        "weight_name": "p_q_image_2.1_5_steps_v0.1.safetensors",
+        "adapter_name": "step5_v01",
+        "steps": 5,
+        "sigmas": [1.0, 0.94, 6 / 7, 2 / 3, 0.4],
+        "dynamic_shifting": False,
+        "use_kv_cache": True,
+    },
+}
+
+VARIANT_ALIASES = {
+    "6-steps": "6-steps",
+    "6-step": "6-steps",
+    "6step": "6-steps",
+    "6": "6-steps",
+    "basic": "6-steps",
+    "v0.2": "6-steps",
+    "v0.2-6-step": "6-steps",
+    "v0.2-6-steps": "6-steps",
+    "6-steps-upsampled": "6-steps-upsampled",
+    "6-step-upsampled": "6-steps-upsampled",
+    "6step-upsampled": "6-steps-upsampled",
+    "prompt-upsampled": "6-steps-upsampled",
+    "upsampled": "6-steps-upsampled",
+    "upsample": "6-steps-upsampled",
+    "v0.2-upsampled": "6-steps-upsampled",
+    "v0.2-6-step-upsampled": "6-steps-upsampled",
+    "v0.2-6-steps-upsampled": "6-steps-upsampled",
+    "8-steps": "8-steps",
+    "8-step": "8-steps",
+    "8step": "8-steps",
+    "8": "8-steps",
+    "v0.1-8-step": "8-steps",
+    "v0.1-8-steps": "8-steps",
+    "5-steps": "5-steps",
+    "5-step": "5-steps",
+    "5step": "5-steps",
+    "5": "5-steps",
+    "v0.1-5-step": "5-steps",
+    "v0.1-5-steps": "5-steps",
 }
 
-# Load the base pipeline once at module level (ZeroGPU CUDA emulation makes
-# this a no-op until a @spaces.GPU function actually runs).
 pipe = QwenImage21Pipeline.from_pretrained(
     "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
-).to("cuda")
+)
+for spec in VARIANTS.values():
+    pipe.load_lora_weights(
+        LORA_REPO,
+        weight_name=spec["weight_name"],
+        adapter_name=spec["adapter_name"],
+    )
+pipe.set_adapters([VARIANTS[DEFAULT_VARIANT]["adapter_name"]], [LORA_STRENGTH])
 pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
     pipe.scheduler.config,
-    use_dynamic_shifting=False,
-    shift=1.0,
+    use_dynamic_shifting=True,
     shift_terminal=None,
 )
+pipe = pipe.to("cuda")
 
-_loaded_steps = None
+_active_variant = DEFAULT_VARIANT
 
 
-def _ensure_adapter(steps: int):
-    """Load exactly one Pruna LoRA adapter matching the requested step count."""
-    global _loaded_steps
-    steps = 5 if steps is not None and int(steps) == 5 else 8
-    if _loaded_steps != steps:
-        pipe.unload_lora_weights()
-        pipe.load_lora_weights(
-            "PrunaAI/Pruna-Qwen-Image-2.1",
-            weight_name=f"p_qwen_image_2.1_{steps}step_v0.1.safetensors",
+def _resolve_variant(variant) -> str:
+    if variant is None or variant == "":
+        return DEFAULT_VARIANT
+    key = "-".join(str(variant).strip().lower().replace("_", "-").split())
+    try:
+        key = str(int(float(key)))
+    except (TypeError, ValueError):
+        pass
+    return VARIANT_ALIASES.get(key, DEFAULT_VARIANT)
+
+
+def _ensure_variant(variant) -> dict:
+    """Activate the matching adapter and scheduler. Load only one at a time."""
+    global _active_variant
+    name = _resolve_variant(variant)
+    spec = VARIANTS[name]
+    if _active_variant != name:
+        pipe.set_adapters([spec["adapter_name"]], [LORA_STRENGTH])
+        scheduler_kwargs = {"use_dynamic_shifting": spec["dynamic_shifting"], "shift_terminal": None}
+        if not spec["dynamic_shifting"]:
+            scheduler_kwargs["shift"] = 1.0
+        pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+            pipe.scheduler.config,
+            **scheduler_kwargs,
         )
-        _loaded_steps = steps
-    return steps
+        _active_variant = name
+    return spec
 
 
 def _seed(seed) -> int:
@@ -78,37 +165,36 @@ def _to_file(img: Image.Image) -> dict:
     return {"path": path, "url": url, "is_file": True}
 
 
+def _run_pipe(prompt: str, spec: dict, seed: int, image=None):
+    kwargs = {
+        "prompt": prompt,
+        "generator": torch.Generator("cuda").manual_seed(_seed(seed)),
+        "num_inference_steps": spec["steps"],
+        "sigmas": spec["sigmas"],
+        "true_cfg_scale": 1.0,
+        "use_kv_cache": spec["use_kv_cache"],
+        "output_resolution": 1024,
+    }
+    if image is None:
+        kwargs["width"] = 1024
+        kwargs["height"] = 1024
+    else:
+        kwargs["image"] = image
+    return pipe(**kwargs).images[0]
+
+
 @spaces.GPU(duration=120)
-def generate(prompt: str, steps: int = 8, seed: int = 42) -> dict:
-    """Text-to-image with Pruna-Qwen-Image-2.1 (5 or 8 steps, no CFG)."""
-    steps = _ensure_adapter(steps)
-    out = pipe(
-        prompt=prompt,
-        width=1024,
-        height=1024,
-        generator=torch.Generator("cuda").manual_seed(_seed(seed)),
-        num_inference_steps=steps,
-        sigmas=SIGMAS[steps],
-        true_cfg_scale=1.0,
-        use_kv_cache=True,
-    ).images[0]
-    return _to_file(out)
+def generate(prompt: str, variant: str = DEFAULT_VARIANT, seed: int = 42) -> dict:
+    """Text-to-image with P-Q-Image-2.1 (no CFG)."""
+    spec = _ensure_variant(variant)
+    return _to_file(_run_pipe(prompt, spec, seed))
 
 
 @spaces.GPU(duration=120)
-def edit(image, prompt: str, steps: int = 8, seed: int = 42) -> dict:
-    """Edit an image with Pruna-Qwen-Image-2.1 (5 or 8 steps, no CFG)."""
-    steps = _ensure_adapter(steps)
-    out = pipe(
-        prompt=prompt,
-        image=_coerce_image(image),
-        generator=torch.Generator("cuda").manual_seed(_seed(seed)),
-        num_inference_steps=steps,
-        sigmas=SIGMAS[steps],
-        true_cfg_scale=1.0,
-        use_kv_cache=True,
-    ).images[0]
-    return _to_file(out)
+def edit(image, prompt: str, variant: str = DEFAULT_VARIANT, seed: int = 42) -> dict:
+    """Edit an image with P-Q-Image-2.1 (no CFG)."""
+    spec = _ensure_variant(variant)
+    return _to_file(_run_pipe(prompt, spec, seed, image=_coerce_image(image)))
 
 
 demo = gr.Workflow(

workflow.json

@@ -1,6 +1,6 @@
 {
   "schema_version": "2",
-  "name": "Pruna Qwen Image 2.1",
+  "name": "P-Q-Image-2.1",
   "references": [
     {
       "id": "ref_prompt",
@@ -14,15 +14,15 @@
       "data": {"out": "A glowing neon shop sign that reads \"QWEN IMAGE 2.1\", mounted on a brick wall in a narrow city alley at night. Heavy rain, wet pavement reflecting pink and blue light, shallow depth of field, cinematic photograph."}
     },
     {
-      "id": "ref_steps",
-      "label": "Steps (5 or 8)",
+      "id": "ref_variant",
+      "label": "Variant (6-steps | 6-steps-upsampled | 8-steps | 5-steps)",
       "role": "reference",
-      "asset_type": "number",
+      "asset_type": "text",
       "x": 40,
       "y": 320,
-      "inputs": [{"id": "in", "label": "Number", "type": "number"}],
-      "outputs": [{"id": "out", "label": "Number", "type": "number"}],
-      "data": {"out": 8}
+      "inputs": [{"id": "in", "label": "Text", "type": "text"}],
+      "outputs": [{"id": "out", "label": "Text", "type": "text"}],
+      "data": {"out": "6-steps"}
     },
     {
       "id": "ref_image",
@@ -68,7 +68,7 @@
       "y": 100,
       "inputs": [
         {"id": "prompt", "label": "Prompt", "type": "text", "required": true},
-        {"id": "steps", "label": "Steps", "type": "number"},
+        {"id": "variant", "label": "Variant", "type": "text"},
         {"id": "seed", "label": "Seed", "type": "number"}
       ],
       "outputs": [{"id": "out_0", "label": "Image", "type": "image", "output_index": 0}]
@@ -84,7 +84,7 @@
       "inputs": [
         {"id": "image", "label": "Image", "type": "image", "required": true},
         {"id": "prompt", "label": "Edit Prompt", "type": "text", "required": true},
-        {"id": "steps", "label": "Steps", "type": "number"},
+        {"id": "variant", "label": "Variant", "type": "text"},
         {"id": "seed", "label": "Seed", "type": "number"}
       ],
       "outputs": [{"id": "out_0", "label": "Edited Image", "type": "image", "output_index": 0}]
@@ -114,12 +114,12 @@
   ],
   "edges": [
     {"id": "e1", "from_node_id": "ref_prompt", "from_port_id": "out", "to_node_id": "op_generate", "to_port_id": "prompt", "type": "text"},
-    {"id": "e2", "from_node_id": "ref_steps", "from_port_id": "out", "to_node_id": "op_generate", "to_port_id": "steps", "type": "number"},
+    {"id": "e2", "from_node_id": "ref_variant", "from_port_id": "out", "to_node_id": "op_generate", "to_port_id": "variant", "type": "text"},
     {"id": "e3", "from_node_id": "ref_seed", "from_port_id": "out", "to_node_id": "op_generate", "to_port_id": "seed", "type": "number"},
     {"id": "e4", "from_node_id": "ref_image", "from_port_id": "out", "to_node_id": "op_edit", "to_port_id": "image", "type": "image"},
     {"id": "e7", "from_node_id": "ref_edit_prompt", "from_port_id": "out", "to_node_id": "op_edit", "to_port_id": "prompt", "type": "text"},
     {"id": "e5", "from_node_id": "op_generate", "from_port_id": "out_0", "to_node_id": "sub_generated", "to_port_id": "in", "type": "image"},
-    {"id": "e8", "from_node_id": "ref_steps", "from_port_id": "out", "to_node_id": "op_edit", "to_port_id": "steps", "type": "number"},
+    {"id": "e8", "from_node_id": "ref_variant", "from_port_id": "out", "to_node_id": "op_edit", "to_port_id": "variant", "type": "text"},
     {"id": "e9", "from_node_id": "ref_seed", "from_port_id": "out", "to_node_id": "op_edit", "to_port_id": "seed", "type": "number"},
     {"id": "e6", "from_node_id": "op_edit", "from_port_id": "out_0", "to_node_id": "sub_edited", "to_port_id": "in", "type": "image"}
   ]