import gradio as gr import math import numpy as np import random import torch import spaces import os import requests import tempfile import shutil from PIL import Image from diffusers import QwenImageEditPlusPipeline from typing import List, Tuple from urllib.parse import urlparse MAX_SEED = np.iinfo(np.int32).max # --- Model Loading --- dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" pipe = QwenImageEditPlusPipeline.from_pretrained( "Qwen/Qwen-Image-Edit-2511", torch_dtype=dtype ).to(device) # Fuse the lightning LoRA directly into the base weights pipe.load_lora_weights( "lightx2v/Qwen-Image-Edit-2511-Lightning", weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors", ) pipe.fuse_lora() pipe.unload_lora_weights() _VAE_IMAGE_SIZE = 1024 * 1024 def calculate_vae_gen_size(image: Image.Image) -> tuple: W, H = image.size ratio = W / H gen_w = math.sqrt(_VAE_IMAGE_SIZE * ratio) gen_h = gen_w / ratio gen_w = round(gen_w / 32) * 32 gen_h = round(gen_h / 32) * 32 return int(gen_w), int(gen_h) def resize_image(image: Image.Image) -> Image.Image: """Cap longest side to 1328px, snap to multiples of 16.""" MAX_SIDE = 1328 w, h = image.size scale = min(MAX_SIDE / w, MAX_SIDE / h, 1.0) new_w = (int(w * scale) // 16) * 16 new_h = (int(h * scale) // 16) * 16 if (new_w, new_h) == (w, h): return image return image.resize((new_w, new_h), Image.LANCZOS) def load_lora_auto(pipe, lora_input: str): """Load LoRA from HuggingFace repo ID, URL, or blob link.""" lora_input = lora_input.strip() if not lora_input: return False if "/" in lora_input and not lora_input.startswith("http"): pipe.load_lora_weights(lora_input) return True if lora_input.startswith("http"): url = lora_input if "huggingface.co" in url and "/blob/" not in url and "/resolve/" not in url: repo_id = urlparse(url).path.strip("/") pipe.load_lora_weights(repo_id) return True if "/blob/" in url: url = url.replace("/blob/", "/resolve/") tmp_dir = tempfile.mkdtemp() local_path = os.path.join(tmp_dir, os.path.basename(urlparse(url).path)) try: print(f"Downloading LoRA from {url}...") resp = requests.get(url, stream=True) resp.raise_for_status() with open(local_path, "wb") as f: for chunk in resp.iter_content(chunk_size=8192): f.write(chunk) pipe.load_lora_weights(local_path) return True finally: shutil.rmtree(tmp_dir, ignore_errors=True) return False @spaces.GPU def infer( gallery_images, prompt: str, lora_id: str = "", seed: int = 0, randomize_seed: bool = True, true_guidance_scale: float = 1.0, num_inference_steps: int = 4, width: int = 1024, height: int = 1024, auto_size: bool = True, progress=gr.Progress(track_tqdm=True) ) -> Tuple[Image.Image, int]: # gallery_images is a list of (pil_image, caption) tuples or just pil images if not gallery_images: raise gr.Error("Please upload at least 1 image.") # images = [resize_image(img[0] if isinstance(img, tuple) else img).convert("RGB") # for img in gallery_images[:3]] processed_images = [] for item in gallery_images[:3]: # Gradio gallery yields (image, caption) tuples or dictionaries depending on version img_obj = item[0] if isinstance(item, tuple) else (item.image if hasattr(item, 'image') else item) # Apply your image scaling constraints and convert to RGB processed_images.append(resize_image(img_obj).convert("RGB")) images = processed_images if len(gallery_images) > 3: gr.Warning("Only the first 3 images are used.") if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) print(f"Running with {len(images)} input image(s), prompt: {prompt!r}") custom_lora_loaded = False if lora_id and lora_id.strip(): try: custom_lora_loaded = load_lora_auto(pipe, lora_id) print(f"Loaded custom LoRA: {lora_id}") except Exception as e: print(f"LoRA load failed: {e}") custom_lora_loaded = False if auto_size: width, height = calculate_vae_gen_size(images[0]) try: result = pipe( image=images, prompt=prompt, height=height, width=width, num_inference_steps=num_inference_steps, generator=generator, true_cfg_scale=true_guidance_scale, num_images_per_prompt=1, ).images[0] finally: if custom_lora_loaded: pipe.unload_lora_weights() return result, seed # # --- UI --- # css = "#col-container { max-width: 1100px; margin: 0 auto; }" # --- UI --- css = ''' #col-container { max-width: 1000px; margin: 0 auto; } .dark .progress-text { color: white !important } #examples { max-width: 1000px; margin: 0 auto; } .image-container { min-height: 300px; } ''' POPULAR_LORAS = [ ("🎨 Color Grade Transfer", "ovi054/QIE-2511-Color-Grade-Transfer-LoRA"), ] with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown("## 🖌️ Qwen Image Edit 2511 — LoRA Studio") gr.Markdown( "Upload **1–3 images** via the gallery, write a prompt, and optionally apply a LoRA. " "The order of images in the gallery is Image 1, 2, 3." ) with gr.Row(): with gr.Column(): input_gallery = gr.Gallery( label="Input Images (upload 1–3)", columns=3, rows=1, # height=300, object_fit="contain", type="pil", interactive=True, ) prompt = gr.Textbox( label="Prompt", # placeholder="e.g. 'Put the person from Image 1 into the scene from Image 2'", lines=2, ) lora_id = gr.Textbox( label="LoRA (repo ID or URL)", placeholder="author/model or https://huggingface.co/…/model.safetensors", ) run_btn = gr.Button("🎨 Run Edit", variant="primary", size="lg") with gr.Accordion("⚙️ Advanced Settings", open=False): with gr.Row(): width = gr.Slider(label="Width", value=1024, minimum=64, maximum=2048, step=16) height = gr.Slider(label="Height", value=1024, minimum=64, maximum=2048, step=16) auto_size = gr.Checkbox(label="Auto size", value=True) seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) true_guidance_scale = gr.Slider( label="True Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0 ) num_inference_steps = gr.Slider( label="Inference Steps", minimum=1, maximum=40, step=1, value=4 ) with gr.Column(): result = gr.Image(label="✨ Output", interactive=False) gr.Markdown("**Quick LoRAs:**") with gr.Row(): for btn_label, repo in POPULAR_LORAS: gr.Button(btn_label, size="sm", variant="secondary").click( fn=lambda r=repo: r, outputs=[lora_id] ) # output_seed = gr.Number(label="Seed used", precision=0) run_btn.click( fn=infer, inputs=[input_gallery, prompt, lora_id, seed, randomize_seed, true_guidance_scale, num_inference_steps, width, height, auto_size], outputs=[result, seed] ) demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=css)