"""gr.Workflow app for PrunaAI/P-Q-Image-2.1 (ZeroGPU). Nodes on the canvas: - "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 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 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_CHOICES = list(VARIANTS) 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", } pipe = QwenImage21Pipeline.from_pretrained( "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16 ) 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=True, shift_terminal=None, ) pipe = pipe.to("cuda") _active_variant = DEFAULT_VARIANT 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, ) _active_variant = name return spec def _seed(seed) -> int: try: return int(seed) except (TypeError, ValueError): return 42 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} 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] def _variant_dropdown() -> gr.Dropdown: return gr.Dropdown( choices=VARIANT_CHOICES, value=DEFAULT_VARIANT, label="Variant", ) @spaces.GPU(duration=120) def generate( prompt: str, variant: str = _variant_dropdown(), 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, variant: str = _variant_dropdown(), 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( graph="workflow.json", bind={"Generate Image": generate, "Edit Image": edit}, ) if __name__ == "__main__": demo.launch()