""" Pruna-Qwen-Image-2.1 — a Gradio demo Space for the PrunaAI DMD-distillation LoRA on Qwen/Qwen-Image-2.1. One pipeline does both tasks, exactly as the base does: leave the reference gallery empty for text-to-image, or attach up to three reference images for editing / composition / style transfer. The distillation recipe follows the Pruna model card 1:1: - 5-step and 8-step adapters, each trained on its own sigma schedule (raw nodes passed as `sigmas=`). - Scheduler rebuilt with use_dynamic_shifting=False, shift=1.0, shift_terminal=None so the given sigmas are used verbatim (the stock config would shift them a second time and stretch the terminal node to 0.02). - true_cfg_scale=1.0, no negative prompt, LoRA strength 1.0, use_kv_cache=True. - Condition images encoded at 1024-area (output_resolution=1024), the training condition resolution. """ import os # Before torch: ZeroGPU. expandable_segments avoids the fragmentation OOM that shows up on ZeroGPU # as "NVML_SUCCESS == r INTERNAL ASSERT FAILED" from the caching allocator. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 # must precede any torch import on ZeroGPU import math # noqa: E402 import random # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 from pathlib import Path # noqa: E402 import gradio as gr # noqa: E402 import torch # noqa: E402 from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline # noqa: E402 from PIL import Image # noqa: E402 # -------------------------------------------------------------------------------------- # Constants # -------------------------------------------------------------------------------------- BASE_MODEL_ID = "Qwen/Qwen-Image-2.1" LORA_REPO = "PrunaAI/Pruna-Qwen-Image-2.1" # Each adapter is trained on its own sigma schedule; only one is active on the transformer at a # time (switched with set_adapters). The terminal 0 is appended by the scheduler. STEPS_TO_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], } ADAPTER_NAMES = {5: "pruna_5step", 8: "pruna_8step"} DEFAULT_STEPS = 8 # the card's recommended default: 8 for quality, 5 for speed LORA_STRENGTH = 1.0 MAX_SEED = 2**31 - 1 MAX_REFS = 3 # the card supports single- and multi-image editing with up to 3 references EXAMPLES_DIR = Path(__file__).parent / "examples" # 1024-area output buckets (the distillation training resolution) across common aspect ratios. # calculate_dimensions is the pipeline's own rule: sqrt-area sizing rounded to a multiple of 32. ASPECT_RATIOS = { "1:1": 1.0, "16:9": 16 / 9, "9:16": 9 / 16, "4:3": 4 / 3, "3:4": 3 / 4, "3:2": 3 / 2, "2:3": 2 / 3, } TARGET_AREA = 1024 * 1024 def calculate_dimensions(target_area, ratio): width = math.sqrt(target_area * ratio) height = width / ratio return round(width / 32) * 32, round(height / 32) * 32 SIZE_CHOICES = [] LABEL_TO_DIMS = {} for _label, _ratio in ASPECT_RATIOS.items(): _w, _h = calculate_dimensions(TARGET_AREA, _ratio) _choice = f"{_label} · {_w}×{_h}" SIZE_CHOICES.append(_choice) LABEL_TO_DIMS[_choice] = (_w, _h) # -------------------------------------------------------------------------------------- # Model load (module scope, eagerly on cuda) # -------------------------------------------------------------------------------------- print("Loading base pipeline…", flush=True) pipe = QwenImage21Pipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16) # Both adapters are fetched at startup (336 MB each) and attached under distinct PEFT names; # set_adapters activates exactly one, honouring the card's "load one adapter at a time" rule — # the transformer never sees both schedules' deltas simultaneously. for _steps, _adapter_name in ADAPTER_NAMES.items(): pipe.load_lora_weights( LORA_REPO, weight_name=f"p_qwen_image_2.1_{_steps}step_v0.1.safetensors", adapter_name=_adapter_name, ) pipe.set_adapters([ADAPTER_NAMES[DEFAULT_STEPS]], [LORA_STRENGTH]) # Use the card's sigma nodes verbatim: no extra shifting. (The stock scheduler config has # use_dynamic_shifting=true and shift_terminal=0.02, which would alter the passed sigmas.) pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config( pipe.scheduler.config, use_dynamic_shifting=False, shift=1.0, shift_terminal=None, ) pipe = pipe.to("cuda") # transformers' Qwen3VLVisionPatchEmbed runs an nn.Conv3d whose kernel_size equals its stride — # exactly a linear map over each flattened patch. cuDNN has no usable bf16 Conv3d kernel for this # shape; the conv costs ~30 s per reference image in this pipeline. The matmul below is 0.8 ms and # bitwise equal to the conv's output. patch_embed = pipe.text_encoder.model.visual.patch_embed def _patch_embed_forward(hidden_states): proj = patch_embed.proj weight = proj.weight.reshape(proj.weight.shape[0], -1) return hidden_states.to(weight.dtype).flatten(1) @ weight.T + proj.bias patch_embed.forward = _patch_embed_forward print("Pipeline ready.", flush=True) # -------------------------------------------------------------------------------------- # Inference # -------------------------------------------------------------------------------------- def normalize_gallery(gallery) -> list[str]: """Reduce a ``gr.Gallery`` value to a flat list of local file paths.""" paths: list[str] = [] for item in gallery or []: if isinstance(item, (tuple, list)): item = item[0] if isinstance(item, dict): item = item.get("path") or item.get("image") or item.get("name") if isinstance(item, dict): item = item.get("path") if isinstance(item, Image.Image): tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False) item.save(tmp.name) item = tmp.name if isinstance(item, str) and item: paths.append(item) return paths # Measured on this Space: 8-step T2I 4.5 s, 1-ref edit 4.7 s, 5-step 2.4 s, worst supported case # (3-ref edit at 1376×768) 5.2 s on the GPU. 20 s leaves ~4× headroom for pool variance without # padding every visitor's quota. size="xlarge" is required: a 3-ref 1024-area edit allocates # ~44 GiB, over the ~44.7 GiB usable on `large`. @spaces.GPU(duration=20, size="xlarge") def generate( prompt, input_images=None, steps=DEFAULT_STEPS, size_label=SIZE_CHOICES[0], seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True), ): prompt = (prompt or "").strip() if not prompt: raise gr.Error("Please enter a prompt.") steps = int(steps) if steps not in STEPS_TO_SIGMAS: raise gr.Error( f"Steps must be one of {sorted(STEPS_TO_SIGMAS)} — each adapter is trained on its own sigma schedule." ) image_paths = normalize_gallery(input_images) if len(image_paths) > MAX_REFS: raise gr.Error(f"Up to {MAX_REFS} reference images are supported.") refs = [Image.open(p).convert("RGB") for p in image_paths] or None width, height = LABEL_TO_DIMS.get(size_label, LABEL_TO_DIMS[SIZE_CHOICES[0]]) seed = int(seed) if seed is not None else 0 if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator("cuda").manual_seed(seed) # Activate exactly the adapter matching the chosen schedule. pipe.set_adapters([ADAPTER_NAMES[steps]], [LORA_STRENGTH]) started = time.perf_counter() try: result = pipe( prompt=prompt, image=refs, width=width, height=height, num_inference_steps=steps, sigmas=STEPS_TO_SIGMAS[steps], true_cfg_scale=1.0, # no CFG: the adapters are distilled without it use_kv_cache=True, output_resolution=1024, # condition images encoded at the training resolution generator=generator, ).images[0] except torch.cuda.OutOfMemoryError: raise gr.Error( f"Out of GPU memory with {len(image_paths)} reference image(s) at {width}×{height}. " "Try fewer references or a smaller size." ) finally: torch.cuda.empty_cache() elapsed = time.perf_counter() - started n_refs = len(image_paths) mode = f"edit · {n_refs} reference{'s' if n_refs != 1 else ''}" if n_refs else "text-to-image" info = ( f"{steps} steps · {width}×{height} · {mode} · seed {seed} · {elapsed:.1f}s on the GPU · " f"adapter p_qwen_image_2.1_{steps}step_v0.1" ) return result, seed, info # -------------------------------------------------------------------------------------- # Examples # -------------------------------------------------------------------------------------- # Reference photos baked into the Space for the edit examples — CC0, from # linoyts/repo-to-space-example-inputs. A row's reference cell is a LIST of file paths (one per # reference image), which gr.Examples maps onto the Gallery input; an empty list = no references. REF_WOMAN = str(EXAMPLES_DIR / "woman.jpg") REF_CAMERA = str(EXAMPLES_DIR / "vintage_camera.jpg") REF_CAT = str(EXAMPLES_DIR / "cat_tabby.jpg") # Rows are the positional prefix of generate(): prompt, input_images, steps, size_label, seed, # randomize_seed. randomize_seed=False makes every cached row reproducible — click a row, press # Generate, get the same image. T2I_EXAMPLES = [ # The Pruna model card's own T2I showcase prompt, verbatim. [ '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.", [], 8, SIZE_CHOICES[0], 42, False, ], [ "A close-up photo of a bumblebee perched on a pink zinnia flower at dawn, dew drops on the " "petals, soft golden backlight, shallow depth of field, macro photography.", [], 5, SIZE_CHOICES[0], 42, False, ], ] EDIT_EXAMPLES = [ # The card's own editing example ("Change the background to a sunset beach") on a CC0 portrait. [ "Change the background to a sunset beach", [REF_WOMAN], 8, SIZE_CHOICES[0], 42, False, ], # A restyle edit on the same CC0 portrait. [ "Turn this into a pencil sketch portrait on white paper", [REF_WOMAN], 8, SIZE_CHOICES[0], 42, False, ], # A product-style scene edit with a second CC0 image. [ "Place this camera on a photographer's cluttered desk next to a cup of coffee, warm morning light", [REF_CAMERA], 8, SIZE_CHOICES[0], 42, False, ], # Multi-reference composition: combine two subjects into one scene. [ "Combine these two animals into one image: a cat sitting on sunlit grass, golden afternoon light.", [REF_WOMAN, REF_CAT], 8, SIZE_CHOICES[0], 42, False, ], ] # -------------------------------------------------------------------------------------- # UI # -------------------------------------------------------------------------------------- css = """ #col-container { margin: 0 auto; max-width: 1100px; } /* Citrus leaves some text colours unset in dark mode, which renders dark-on-dark. */ .dark .gradio-container { color: var(--body-text-color); } """ # `theme` / `css` are passed to launch(), not to the Blocks constructor: Gradio 6 moved them and # warns if they are set here. with gr.Blocks(title="Pruna-Qwen-Image-2.1") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # ⚡ Pruna-Qwen-Image-2.1 **5 / 8-step Qwen-Image-2.1.** A DMD-distillation LoRA from [PrunaAI](https://huggingface.co/PrunaAI/Pruna-Qwen-Image-2.1) that runs [Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) in a few steps with no classifier-free guidance — the same pipeline, a fraction of the 40-step time. Leave the reference images empty for **text-to-image**; attach up to three for **editing / composition / style transfer**. """ ) with gr.Row(): with gr.Column(): prompt = gr.Textbox( label="Prompt", placeholder=( "Describe the image to generate — detailed prompts work much better. " "Or an edit instruction if you attach reference images." ), lines=4, ) input_images = gr.Gallery( label="Reference images (editing, up to 3 — leave empty for text-to-image)", type="filepath", interactive=True, columns=3, height=220, ) with gr.Accordion("Advanced options", open=False): steps = gr.Radio( choices=[8, 5], value=DEFAULT_STEPS, label="Sampling steps (adapter)", info="8 steps: higher quality (recommended). 5 steps: faster, visibly lower quality. " "Each adapter is trained on its own sigma schedule; switching activates it.", ) size_label = gr.Dropdown( choices=SIZE_CHOICES, value=SIZE_CHOICES[0], label="Output size (1024-area, the training resolution)", ) seed = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=MAX_SEED) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) generate_btn = gr.Button("Generate", variant="primary") with gr.Column(): output_image = gr.Image(label="Result", type="pil", height=520) used_seed = gr.Number(label="Seed used", precision=0, interactive=False) info = gr.Markdown() gr.Markdown("**Text-to-image** — leave the references empty; detailed prompts work much better") gr.Examples( examples=T2I_EXAMPLES, inputs=[prompt, input_images, steps, size_label, seed, randomize_seed], fn=generate, outputs=[output_image, used_seed, info], cache_examples=True, cache_mode="lazy", examples_per_page=8, label="Text-to-image cases — click a row to run it; its steps / size / seed load into the form", ) gr.Markdown("**Image editing** — each row loads its reference images too") gr.Examples( examples=EDIT_EXAMPLES, inputs=[prompt, input_images, steps, size_label, seed, randomize_seed], fn=generate, outputs=[output_image, used_seed, info], cache_examples=True, cache_mode="lazy", examples_per_page=8, label="Editing cases — click a row to run it; its references / steps / size / seed load into the form", ) gr.Markdown( "---\n" "Model: [PrunaAI/Pruna-Qwen-Image-2.1](https://huggingface.co/PrunaAI/Pruna-Qwen-Image-2.1) · " "distilled from [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1), " "released under the **Qwen RESEARCH LICENSE AGREEMENT** (non-commercial, research/evaluation only). " "v0.1: quality is below the base model — for the full 40-step quality use the " "[base Space](https://huggingface.co/spaces/hugging-apps/qwen-image-2-1). " "Example reference photos are CC0 from " "[linoyts/repo-to-space-example-inputs](https://huggingface.co/datasets/linoyts/repo-to-space-example-inputs)." ) generate_btn.click( fn=generate, inputs=[prompt, input_images, steps, size_label, seed, randomize_seed], outputs=[output_image, used_seed, info], ) demo.queue(default_concurrency_limit=1, max_size=25) if __name__ == "__main__": demo.launch( mcp_server=True, show_error=True, theme=gr.themes.Citrus(), css=css, allowed_paths=[str(EXAMPLES_DIR)], )