Spaces:
Running on Zero
Running on Zero
| """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 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", | |
| ) | |
| _loaded_steps = steps | |
| return steps | |
| 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 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 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) | |
| demo = gr.Workflow( | |
| graph="workflow.json", | |
| bind={"Generate Image": generate, "Edit Image": edit}, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |