Add gr.Workflow app for Pruna-Qwen-Image-2.1 (ZeroGPU, gradio 6.28.0)

README.md

@@ -4,8 +4,8 @@ emoji: ⚡
 colorFrom: indigo
 colorTo: indigo
 sdk: gradio
-sdk_version: 6.25.0
-app_file: run.py
+sdk_version: 6.28.0
+app_file: app.py
 pinned: false
 hf_oauth: true
 hf_oauth_scopes:

app.py

@@ -0,0 +1,113 @@
+"""gr.Workflow app for PrunaAI/Pruna-Qwen-Image-2.1 (ZeroGPU).
+
+Nodes on the canvas:
+  - "Generate Image": text-to-image with the 5- or 8-step LoRA adapter
+  - "Edit Image": image editing with an input image + prompt
+"""
+
+import os
+
+import torch
+import spaces
+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],
+}
+
+# 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")
+pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+    pipe.scheduler.config,
+    use_dynamic_shifting=False,
+    shift=1.0,
+    shift_terminal=None,
+)
+
+_loaded_steps = None
+
+
+def _ensure_adapter(steps: int):
+    """Load exactly one Pruna LoRA adapter matching the requested step count."""
+    global _loaded_steps
+    steps = 5 if 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",
+        )
+        _loaded_steps = steps
+    return steps
+
+
+def _coerce_image(image):
+    """Accept a PIL image, a filepath, or a Gradio/Workflow file-ref dict."""
+    if image is None or image == "":
+        return None
+    if isinstance(image, Image.Image):
+        return image.convert("RGB")
+    if isinstance(image, dict):
+        image = image.get("path") or image.get("url") or ""
+    return Image.open(str(image)).convert("RGB")
+
+
+def _to_file(img: Image.Image) -> dict:
+    """Save a PIL image into Gradio's upload folder and return a file-ref dict
+    (same shape gr.Workflow uses internally for media node outputs)."""
+    directory = get_upload_folder()
+    os.makedirs(directory, exist_ok=True)
+    path = os.path.join(directory, f"pruna_{os.urandom(8).hex()}.png")
+    img.save(path)
+    url = f"/gradio_api/file={client_utils.encode_file_path(path)}"
+    return {"path": path, "url": url, "is_file": True}
+
+
+@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(int(seed)),
+        num_inference_steps=steps,
+        sigmas=SIGMAS[steps],
+        true_cfg_scale=1.0,
+        use_kv_cache=True,
+    ).images[0]
+    return _to_file(out)
+
+
+@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(int(seed)),
+        num_inference_steps=steps,
+        sigmas=SIGMAS[steps],
+        true_cfg_scale=1.0,
+        use_kv_cache=True,
+    ).images[0]
+    return _to_file(out)
+
+
+demo = gr.Workflow(
+    graph="workflow.json",
+    bind={"Generate Image": generate, "Edit Image": edit},
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.4.0
+transformers>=5.17
+accelerate
+peft
+pillow
+spaces
+git+https://github.com/huggingface/diffusers@6256aa7666cedd47443adc8f82da9a10e110b09c

run.py

@@ -1,6 +0,0 @@
-import gradio as gr
-
-demo = gr.Workflow()
-
-if __name__ == "__main__":
-    demo.launch()

workflow.json

@@ -1 +1,102 @@
-{"schema_version":"2","name":"My Workflow","runtime":{"default":"client"},"references":[],"operators":[],"subjects":[],"edges":[],"view":{"default":"canvas"}}
+{
+  "schema_version": "2",
+  "name": "Pruna Qwen Image 2.1",
+  "references": [
+    {
+      "id": "ref_prompt",
+      "label": "Prompt",
+      "role": "reference",
+      "asset_type": "text",
+      "x": 40,
+      "y": 80,
+      "inputs": [{"id": "in", "label": "Text", "type": "text"}],
+      "outputs": [{"id": "out", "label": "Text", "type": "text"}],
+      "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)",
+      "role": "reference",
+      "asset_type": "number",
+      "x": 40,
+      "y": 320,
+      "inputs": [{"id": "in", "label": "Number", "type": "number"}],
+      "outputs": [{"id": "out", "label": "Number", "type": "number"}],
+      "data": {"out": 8}
+    },
+    {
+      "id": "ref_seed",
+      "label": "Seed",
+      "role": "reference",
+      "asset_type": "number",
+      "x": 40,
+      "y": 480,
+      "inputs": [{"id": "in", "label": "Number", "type": "number"}],
+      "outputs": [{"id": "out", "label": "Number", "type": "number"}],
+      "data": {"out": 42}
+    }
+  ],
+  "operators": [
+    {
+      "id": "op_generate",
+      "label": "Generate Image",
+      "role": "operator",
+      "kind": "fn",
+      "fn": "Generate Image",
+      "x": 460,
+      "y": 100,
+      "inputs": [
+        {"id": "prompt", "label": "Prompt", "type": "text", "required": true},
+        {"id": "steps", "label": "Steps", "type": "number"},
+        {"id": "seed", "label": "Seed", "type": "number"}
+      ],
+      "outputs": [{"id": "out_0", "label": "Image", "type": "image", "output_index": 0}]
+    },
+    {
+      "id": "op_edit",
+      "label": "Edit Image",
+      "role": "operator",
+      "kind": "fn",
+      "fn": "Edit Image",
+      "x": 860,
+      "y": 120,
+      "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": "seed", "label": "Seed", "type": "number"}
+      ],
+      "outputs": [{"id": "out_0", "label": "Edited Image", "type": "image", "output_index": 0}]
+    }
+  ],
+  "subjects": [
+    {
+      "id": "sub_generated",
+      "label": "Generated Image",
+      "role": "subject",
+      "asset_type": "image",
+      "x": 1260,
+      "y": 80,
+      "inputs": [{"id": "in", "label": "Image", "type": "image"}],
+      "outputs": [{"id": "out", "label": "Image", "type": "image"}]
+    },
+    {
+      "id": "sub_edited",
+      "label": "Edited Image",
+      "role": "subject",
+      "asset_type": "image",
+      "x": 1260,
+      "y": 380,
+      "inputs": [{"id": "in", "label": "Image", "type": "image"}],
+      "outputs": [{"id": "out", "label": "Image", "type": "image"}]
+    }
+  ],
+  "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": "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": "op_generate", "from_port_id": "out_0", "to_node_id": "op_edit", "to_port_id": "image", "type": "image"},
+    {"id": "e5", "from_node_id": "op_generate", "from_port_id": "out_0", "to_node_id": "sub_generated", "to_port_id": "in", "type": "image"},
+    {"id": "e6", "from_node_id": "op_edit", "from_port_id": "out_0", "to_node_id": "sub_edited", "to_port_id": "in", "type": "image"}
+  ]
+}