import spaces import gradio as gr import math import numpy as np import random import torch import os import gc import requests import tempfile import shutil from PIL import Image from diffusers import QwenImageEditPlusPipeline, QwenImageTransformer2DModel from typing import List, Tuple from urllib.parse import urlparse MAX_SEED = np.iinfo(np.int32).max # --- Speed Optimizations for NVIDIA A100 (xlarge) --- torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # --- Model Loading --- dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" pipe = QwenImageEditPlusPipeline.from_pretrained( "Qwen/Qwen-Image-Edit-2511", transformer=QwenImageTransformer2DModel.from_pretrained( "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19", torch_dtype=dtype, device_map="cuda", ), torch_dtype=dtype, ).to(device) _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_single(pipe, lora_input: str, adapter_name: str) -> bool: """Safely fetch and load a single LoRA adapter into diffusers.""" 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, adapter_name=adapter_name) 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, adapter_name=adapter_name) 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, timeout=30) 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, adapter_name=adapter_name) return True finally: shutil.rmtree(tmp_dir, ignore_errors=True) return False @spaces.GPU(size="xlarge") def infer( gallery_images, prompt: str, lora_1: str = "", scale_1: float = 1.0, lora_2: str = "", scale_2: float = 1.0, lora_3: str = "", scale_3: float = 1.0, 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, secret_code: str = "", progress=gr.Progress(track_tqdm=True) ) -> Tuple[Image.Image, int]: # --- Secret Pass Check --- EXPECTED_SECRET = os.getenv("SECRET_CODE") if not EXPECTED_SECRET or secret_code != EXPECTED_SECRET: raise gr.Error("Space is currently unavailable.") gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() if not gallery_images: raise gr.Error("Please upload at least 1 image.") processed_images = [] for item in gallery_images[:3]: img_obj = item[0] if isinstance(item, tuple) else (item.image if hasattr(item, 'image') else item) 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}") # --- Multi-LoRA Loading & Weighting --- active_adapter_names = [] active_adapter_weights = [] lora_requests = [(lora_1, scale_1), (lora_2, scale_2), (lora_3, scale_3)] try: for idx, (lora_input, scale) in enumerate(lora_requests): if lora_input and lora_input.strip(): adapter_name = f"custom_lora_{idx+1}" try: loaded = load_lora_single(pipe, lora_input, adapter_name) # Verify PEFT actually registered the adapter key for Qwen peft_pipe = getattr(pipe, "peft_config", {}) peft_trans = getattr(getattr(pipe, "transformer", None), "peft_config", {}) is_present = (adapter_name in peft_pipe) or (adapter_name in peft_trans) if loaded and is_present: active_adapter_names.append(adapter_name) active_adapter_weights.append(scale) print(f"Successfully loaded LoRA {idx+1}: {lora_input} (weight: {scale})") else: print(f"Skipped LoRA {idx+1} ({lora_input}): Incompatible with Qwen architecture") except Exception as e: print(f"Failed to load LoRA {idx+1} ({lora_input}): {e}") if active_adapter_names: pipe.set_adapters(active_adapter_names, adapter_weights=active_adapter_weights) if auto_size: width, height = calculate_vae_gen_size(images[0]) # --- Fast Inference Execution --- with torch.inference_mode(): result = pipe( image=images, prompt=prompt, negative_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: # Always unload custom LoRAs to ensure 0 VRAM accumulation if active_adapter_names: try: pipe.unload_lora_weights() except Exception as e: print(f"Error unloading LoRAs: {e}") gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() return result, seed # --- UI Setup --- css = ''' #col-container { max-width: 1000px; margin: 0 auto; } .dark .progress-text { color: white !important } .quick-lora-link { display: inline-flex; align-items: center; justify-content: center; width: 28px; height: 28px; border-radius: 6px; background: rgba(128,128,128,0.12); color: inherit; text-decoration: none; font-size: 14px; } ''' POPULAR_LORAS = [ ("🌓 Color Grade Transfer", "ovi054/QIE-2511-Color-Grade-Transfer-LoRA", "Transfer ONLY the color grading from Image 2 onto Image 1", "https://huggingface.co/ovi054/QIE-2511-Color-Grade-Transfer-LoRA"), ("📐 Multiple Angles", "fal/Qwen-Image-Edit-2511-Multiple-Angles-LoRA", " front-left quarter view elevated shot medium shot", "https://huggingface.co/fal/Qwen-Image-Edit-2511-Multiple-Angles-LoRA"), ("✨ Unblur Upscale", "https://huggingface.co/prithivMLmods/Qwen-Image-Edit-2511-Unblur-Upscale/resolve/main/Qwen-Image-Edit-Unblur-Upscale_20.safetensors", "unblur and upscale", "https://huggingface.co/prithivMLmods/Qwen-Image-Edit-2511-Unblur-Upscale"), ("🖼️ Draw2Photo", "ovi054/QIE-2511-Draw2Photo-LoRA", "make it real", "https://huggingface.co/ovi054/QIE-2511-Draw2Photo-LoRA"), ] with gr.Blocks(theme=gr.themes.Citrus(), css=css, analytics_enabled=False) as demo: with gr.Column(elem_id="col-container"): gr.Markdown("## 🖌️ Qwen Image Edit 2511 + Multi-LoRA") with gr.Row(): with gr.Column(): input_gallery = gr.Gallery( label="Input Images (upload 1–3)", columns=3, rows=1, object_fit="contain", type="pil", interactive=True, ) prompt = gr.Textbox( label="Prompt", lines=2, ) with gr.Accordion("🎨 Custom LoRAs (Up to 3)", open=True): with gr.Row(): lora_1 = gr.Textbox(label="LoRA 1 (URL or Repo)", placeholder="repo/name or https://...") scale_1 = gr.Slider(label="LoRA 1 Intensity", minimum=-2.0, maximum=2.0, value=1.0, step=0.05) with gr.Row(): lora_2 = gr.Textbox(label="LoRA 2 (URL or Repo)", placeholder="repo/name or https://...") scale_2 = gr.Slider(label="LoRA 2 Intensity", minimum=-2.0, maximum=2.0, value=1.0, step=0.05) with gr.Row(): lora_3 = gr.Textbox(label="LoRA 3 (URL or Repo)", placeholder="repo/name or https://...") scale_3 = gr.Slider(label="LoRA 3 Intensity", minimum=-2.0, maximum=2.0, value=1.0, step=0.05) # Hidden secret code input for backend validation secret_code_input = gr.Textbox(visible=False, value="", label="Secret Code") run_btn = gr.Button("🎨 Run Edit", variant="primary", size="lg") with gr.Accordion("⚙️ Advanced Settings", open=False): auto_size = gr.Checkbox(label="Auto size", value=True) 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) 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, trigger, url in POPULAR_LORAS: gr.Button(btn_label, size="sm", variant="secondary").click( fn=lambda r=repo, t=trigger: (r, 1.0, t), outputs=[lora_1, scale_1, prompt] ) run_btn.click( fn=infer, inputs=[ input_gallery, prompt, lora_1, scale_1, lora_2, scale_2, lora_3, scale_3, seed, randomize_seed, true_guidance_scale, num_inference_steps, width, height, auto_size, secret_code_input ], outputs=[result, seed] ) demo.launch(mcp_server=True)