Spaces:
Running on Zero
Running on Zero
| """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", | |
| ) | |
| 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)) | |
| 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() | |