import base64 import hashlib import importlib.util import io import json import os import random import re import tempfile import time from pathlib import Path # Configure memory allocator before heavy operations os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") # Rule 1: import spaces FIRST before importing torch or any CUDA-touching library if importlib.util.find_spec("spaces"): import spaces gpu = spaces.GPU(duration=90, size="xlarge") else: gpu = lambda fn: fn import httpx import gradio as gr import torch from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline from huggingface_hub import hf_hub_download from PIL import Image, ImageOps, PngImagePlugin from safetensors.torch import load_file as safetensors_load_file from rendering_config import ( VIGGLE_V03, VIGGLE_V021, NONE_LORA, T2I_CHOICES, resolve_size, sampling_plan, ) # Official BF16 pipeline, matching the Viggle Space. MODEL_ID = "Qwen/Qwen-Image-2.1" MODEL_REVISION = os.environ.get("QWEN_MODEL_REVISION") or None HF_TOKEN = os.environ.get("HF_TOKEN") MODES = [ "Text to Image", "Multi-Image Reference (1-6 Refs)", "Edit Image (1 Ref)", "Transform & Swap (2 Refs)", ] MAX_SEED = 2**31 - 1 # ============================================================ # Curated All-In-One LoRA Specifications # ============================================================ ADAPTER_SPECS = { VIGGLE_V03: { "repo": "Viggle/Qwen-Image-2.1-viggle-turbo", "weights": "Qwen-Image-2.1-viggle-turbo-v0.3-6step-lora-r256.safetensors", "adapter_name": "viggle_turbo_v03", "default_strength": 1.0, "default_steps": 6, "preset_prompt": "", "category": "Turbo Speed", }, "Anime Consistency": { "repo": "WarmBloodAban/Qwen-Image-2.1-LoRAs", "weights": "Qwen2.1_Anime_consistency.safetensors", "adapter_name": "anime_consistency", "default_strength": 0.85, "default_steps": 28, "preset_prompt": "masterpiece, highly detailed anime illustration, vibrant colors, expressive eyes", "category": "Style", }, "Natural Exposure (Photorealism)": { "repo": "prithivMLmods/Qwen-Image-2.1-Natural-Exposure-LoRA", "weights": "Qwen-Image-2.1-Natural-Exposure-LoRA-4000.safetensors", "adapter_name": "natural_exposure", "default_strength": 0.8, "default_steps": 30, "preset_prompt": "natural daylight exposure, authentic colors, unedited 35mm photograph, soft organic textures", "category": "Style", }, "Viggle Turbo (4-Step Acceleration)": { "repo": "Viggle/Qwen-Image-2.1-viggle-turbo", "weights": "Qwen-Image-2.1-viggle-turbo-4step-lora-r64.safetensors", "adapter_name": "viggle_turbo", "default_strength": 1.0, "default_steps": 4, "preset_prompt": "", "category": "Turbo Speed", }, "Viggle Turbo v0.2.1 (6-Step Acceleration)": { "repo": "Viggle/Qwen-Image-2.1-viggle-turbo", "weights": "Qwen-Image-2.1-viggle-turbo-v0.2.1-6step-lora-r256.safetensors", "adapter_name": "viggle_turbo_v021", "default_strength": 1.0, "default_steps": 6, "preset_prompt": "", "category": "Turbo Speed", }, "Fun-Acc (4-Step Turbo)": { "repo": "alibaba-pai/Qwen-Image-2.1-Fun-Acc-LoRAs", "weights": "models/Qwen-Image-2.1-Fun-Acc-4Step.safetensors", "adapter_name": "fun_acc_turbo", "default_strength": 1.0, "default_steps": 4, "preset_prompt": "", "category": "Turbo Speed", }, "BFS Best Face Swap (2 Images)": { "repo": "Alissonerdx/BFS-Best-Face-Swap", "weights": "bfs_head_v1.1_qwen_2.1.safetensors", "adapter_name": "bfs_faceswap", "default_strength": 1.0, "default_steps": 32, "requires_two_images": True, "image2_label": "Upload Head/Face Donor (Image 2)", "needs_alpha_fix": True, "preset_prompt": "head_swap: start with as the base image, keeping its lighting, environment, and background. remove the head from completely and replace it with the head from , strictly preserving the hair, eye color, nose structure from . copy the direction of the eye, head rotation, micro expressions from , high quality, sharp details, 4k", "category": "Two Images", }, "Anything to Real Photo": { "repo": "WarmBloodAban/Qwen-Image-2.1-LoRAs", "weights": "Qwen2.1_Anything2RealCharacters.safetensors", "adapter_name": "any2real", "default_strength": 1.0, "default_steps": 30, "preset_prompt": "change the picture to a realistic high-definition photograph, authentic skin and materials", "category": "Transform", }, "Detail Enhancer": { "repo": "reverentelusarca/elusarcas-qwen-2.1-detail-enhancer-lora", "weights": "elusarcas-qwen2-1-detailer-v1.safetensors", "adapter_name": "detail_enhancer", "default_strength": 0.85, "default_steps": 30, "preset_prompt": "ultra high detail, intricate textures, sharp focus, refined clarity", "category": "Enhance", }, } LORA_CHOICES = [NONE_LORA] + list(ADAPTER_SPECS.keys()) + ["Custom HuggingFace LoRA..."] # Track dynamically loaded adapters LOADED_ADAPTERS = set() # ============================================================ # Official model initialization: no quantization, offload or VAE tiling. # ============================================================ pipe = QwenImage21Pipeline.from_pretrained( MODEL_ID, revision=MODEL_REVISION, dtype=torch.bfloat16, token=HF_TOKEN, ) BASE_SCHEDULER_CONFIG = dict(pipe.scheduler.config) # Keep the original config for non-Viggle adapters; distilled schedules must not stretch the terminal node. pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(BASE_SCHEDULER_CONFIG, shift_terminal=None) for name in (VIGGLE_V03, VIGGLE_V021): spec = ADAPTER_SPECS[name] pipe.load_lora_weights(spec["repo"], weight_name=spec["weights"], adapter_name=spec["adapter_name"], token=HF_TOKEN) LOADED_ADAPTERS.add(spec["adapter_name"]) pipe.to("cuda") # Qwen3-VL non-overlapping Conv3D patches are equivalent to this much faster linear map. patch_embed = pipe.text_encoder.model.visual.patch_embed def _patch_embed_forward(hidden_states): proj = patch_embed.proj weight = proj.weight.reshape(proj.weight.shape[0], -1) return hidden_states.to(weight.dtype).flatten(1) @ weight.T + proj.bias patch_embed.forward = _patch_embed_forward # Conditional K/V must be refreshed when the hybrid tail switches off the adapters. _transformer_forward = pipe.transformer.forward reextract = [False] def transformer_forward(*args, kv_cache_mode=None, **kwargs): if reextract[0] and kv_cache_mode == "cached": kv_cache_mode, reextract[0] = "extract", False return _transformer_forward(*args, kv_cache_mode=kv_cache_mode, **kwargs) pipe.transformer.forward = transformer_forward def base_tail(pipeline, i, t, kwargs): if i == 6: # after seven student steps, run two official-base steps pipeline.disable_lora() reextract[0] = True return kwargs # ============================================================ # LoRA Adapter Loading Helpers # ============================================================ def _inject_missing_alpha_keys(state_dict: dict) -> dict: bases = {} for k, v in state_dict.items(): if not isinstance(v, torch.Tensor): continue if k.endswith(".lora_down.weight") and v.ndim >= 1: base = k[:-len(".lora_down.weight")] rank = int(v.shape[0]) bases[base] = rank elif k.endswith(".lora_A.weight") and v.ndim >= 1: base = k[:-len(".lora_A.weight")] rank = int(v.shape[0]) bases[base] = rank for base, rank in bases.items(): alpha_tensor = torch.tensor(float(rank), dtype=torch.float32) full_alpha = f"{base}.alpha" if full_alpha not in state_dict: state_dict[full_alpha] = alpha_tensor if base.startswith("diffusion_model."): stripped_base = base[len("diffusion_model."):] stripped_alpha = f"{stripped_base}.alpha" if stripped_alpha not in state_dict: state_dict[stripped_alpha] = alpha_tensor return state_dict def _filter_to_diffusers_lora_keys(state_dict: dict) -> dict: keep_suffixes = ( ".lora_up.weight", ".lora_down.weight", ".lora_mid.weight", ".lora_A.weight", ".lora_B.weight", ".alpha", ".lora_alpha", ) out: dict[str, torch.Tensor] = {} for k, v in state_dict.items(): if not isinstance(v, torch.Tensor): continue if k.endswith(".diff") or k.endswith(".diff_b"): continue if not k.endswith(keep_suffixes): continue if k.endswith(".lora_alpha"): base = k[:-len(".lora_alpha")] k2 = f"{base}.alpha" out[k2] = v.float() if v.dtype != torch.float32 else v continue out[k] = v return out def _duplicate_stripped_prefix_keys(state_dict: dict, prefix: str = "diffusion_model.") -> dict: out = dict(state_dict) for k, v in list(state_dict.items()): if not k.startswith(prefix): continue stripped = k[len(prefix):] if stripped not in out: out[stripped] = v return out def _split_gate_up_keys(state_dict: dict) -> dict: """ ai-toolkit / ostris trains fused 'gate_up' linear layers of shape [24576, rank]. Diffusers QwenImage21Transformer2DModel splits this into two separate layers: - 'gate_layer' (first 12288 rows of B) - 'proj' (last 12288 rows of B) Both gate_layer and proj share the same input projection A of shape [rank, 4096]. """ out = {} for k, v in state_dict.items(): if "img_mlp.gate_up" in k: if k.endswith(".lora_B.weight") or k.endswith(".lora_up.weight"): if isinstance(v, torch.Tensor) and v.ndim == 2 and v.shape[0] == 24576: half = 12288 k_gate = k.replace("img_mlp.gate_up", "img_mlp.gate_layer") k_proj = k.replace("img_mlp.gate_up", "img_mlp.proj") out[k_gate] = v[:half, :].clone() out[k_proj] = v[half:, :].clone() continue elif k.endswith(".lora_A.weight") or k.endswith(".lora_down.weight"): k_gate = k.replace("img_mlp.gate_up", "img_mlp.gate_layer") k_proj = k.replace("img_mlp.gate_up", "img_mlp.proj") out[k_gate] = v.clone() if isinstance(v, torch.Tensor) else v out[k_proj] = v.clone() if isinstance(v, torch.Tensor) else v continue elif k.endswith(".alpha") or k.endswith(".lora_alpha"): k_gate = k.replace("img_mlp.gate_up", "img_mlp.gate_layer") k_proj = k.replace("img_mlp.gate_up", "img_mlp.proj") out[k_gate] = v.clone() if isinstance(v, torch.Tensor) else v out[k_proj] = v.clone() if isinstance(v, torch.Tensor) else v continue out[k] = v return out def _check_lora_compatibility(state_dict: dict, adapter_name: str, weight_name: str): """Detect if a LoRA was trained for the legacy 3072-dim Qwen-Image-Edit rather than 4096-dim Qwen-Image-2.1.""" for k, v in state_dict.items(): if isinstance(v, torch.Tensor) and v.ndim == 2: if 3072 in v.shape and 4096 not in v.shape: raise gr.Error( f"LoRA '{adapter_name}' ({weight_name}) is incompatible with Qwen-Image-2.1: " f"parameter shape {list(v.shape)} indicates it was trained for legacy 3072-dim Qwen-Image-Edit, but Qwen-Image-2.1 requires 4096-dim weights." ) def _load_lora_with_fallback(repo: str, weight_name: str, adapter_name: str, needs_alpha_fix: bool = False): try: local_path = hf_hub_download(repo_id=repo, filename=weight_name) sd = safetensors_load_file(local_path) _check_lora_compatibility(sd, adapter_name, weight_name) if any("img_mlp.gate_up" in k for k in sd.keys()): sd = _split_gate_up_keys(sd) if needs_alpha_fix: sd = _inject_missing_alpha_keys(sd) sd = _filter_to_diffusers_lora_keys(sd) sd = _duplicate_stripped_prefix_keys(sd) if not sd: raise gr.Error(f"No valid LoRA weights found in '{weight_name}'.") pipe.load_lora_weights(sd, adapter_name=adapter_name) except Exception as e: err_str = str(e) if "size mismatch" in err_str or "3072" in err_str: raise gr.Error( f"LoRA '{adapter_name}' ({weight_name}) cannot be loaded due to dimension mismatch: " "it was trained for legacy 3072-dim Qwen-Image-Edit, whereas Qwen-Image-2.1 has 4096 hidden dimensions." ) print(f"Pre-processed LoRA load failed ({e}), attempting direct repo load...", flush=True) pipe.load_lora_weights(repo, weight_name=weight_name, adapter_name=adapter_name) def ensure_adapter_ready(selected_lora: str, custom_repo: str = "", custom_file: str = "") -> tuple[str, float]: if selected_lora == NONE_LORA: return "", 1.0 if selected_lora == "Custom HuggingFace LoRA...": custom_repo = (custom_repo or "").strip() custom_file = (custom_file or "").strip() if not custom_repo or not custom_file: raise gr.Error("Please enter both a Hugging Face Repo ID and LoRA filename for Custom LoRA.") adapter_name = f"custom_{hashlib.md5((custom_repo + custom_file).encode()).hexdigest()[:8]}" if adapter_name not in LOADED_ADAPTERS: print(f"Loading custom LoRA from {custom_repo} / {custom_file}...", flush=True) _load_lora_with_fallback(custom_repo, custom_file, adapter_name, needs_alpha_fix=True) LOADED_ADAPTERS.add(adapter_name) return adapter_name, 1.0 spec = ADAPTER_SPECS.get(selected_lora) if not spec: return "", 1.0 adapter_name = spec["adapter_name"] if adapter_name not in LOADED_ADAPTERS: print(f"Loading LoRA {selected_lora} ({spec['repo']} / {spec['weights']})...", flush=True) _load_lora_with_fallback( spec["repo"], spec["weights"], adapter_name, needs_alpha_fix=spec.get("needs_alpha_fix", False), ) LOADED_ADAPTERS.add(adapter_name) return adapter_name, 1.0 # UI strength is the actual adapter scale, not multiplied twice # ============================================================ # Prompt Enhancement (PE) Setup (Viggle Turbo Style) # ============================================================ PE_MODEL_DEFAULT = os.environ.get("PE_MODEL", "deepseek/deepseek-v4.1-flash") PE_T2I_FILE = Path(__file__).parent / "pe_t2i_system_prompt.txt" PE_I2I_FILE = Path(__file__).parent / "pe_i2i_system_prompt.txt" PE_T2I = PE_T2I_FILE.read_text(encoding="utf-8").strip() if PE_T2I_FILE.exists() else ( "You rewrite a user's image request into a prompt for a text-to-image model. " "Reply with the prompt only: one paragraph of plain English, with nothing before or after it: " "no JSON, no label, no markdown, no quotes around it." ) PE_I2I = PE_I2I_FILE.read_text(encoding="utf-8").strip() if PE_I2I_FILE.exists() else ( "You rewrite a user's image-editing instruction into a clear instruction for an image-editing model. " "The input image(s) come with the request. Reply with the instruction only: one paragraph of plain text." ) PE_MAX_PIXELS = 500_000 PE_AUTO_TOKENS = 30 CJK = re.compile(r"[\u4e00-\u9fff]") KANA_HANGUL = re.compile(r"[\u3040-\u30ff\uac00-\ud7af]") def n_tokens(prompt: str) -> int: """Estimate token count in prompt using Qwen tokenizer or regex fallback.""" try: if hasattr(pipe, "processor") and hasattr(pipe.processor, "tokenizer"): return len(pipe.processor.tokenizer(prompt)["input_ids"]) if hasattr(pipe, "tokenizer") and pipe.tokenizer: return len(pipe.tokenizer(prompt)["input_ids"]) except Exception: pass return len(prompt.split()) + len(CJK.findall(prompt)) def enhance_prompt( prompt: str, images: list[Image.Image] | None = None, ratio_name: str = "1:1", api_key: str | None = None, model_name: str | None = None, base_url: str | None = None, ) -> tuple[str | None, str | None]: """Rewrite prompt using DeepSeek V4.1 Flash via OpenRouter/DeepSeek/OpenAI-compatible API.""" images = [img for img in (images or []) if img is not None] if images: chinese = CJK.search(prompt) and not KANA_HANGUL.search(prompt) user_text = f"{prompt}\n\nDescription language: {'Chinese' if chinese else 'English'}" else: user_text = f"{prompt}\nAspect ratio: {ratio_name}" resolved_key = ( (api_key or "").strip() or os.environ.get("OPENROUTER_API_KEY", "").strip() or os.environ.get("DEEPSEEK_API_KEY", "").strip() or os.environ.get("OPENAI_API_KEY", "").strip() ) if not resolved_key: msg = "No API key provided. Please fill in the API Key in Prompt Enhancer settings, or set OPENROUTER_API_KEY / DEEPSEEK_API_KEY." print(f"[pe] skipped: {msg}", flush=True) return None, msg raw_url = (base_url or "").strip() or os.environ.get("PE_BASE_URL", "").strip() raw_model = (model_name or "").strip() or os.environ.get("PE_MODEL", "").strip() # Determine default endpoint if URL is not provided if not raw_url: if resolved_key.startswith("sk-or-"): resolved_url = "https://openrouter.ai/api/v1/chat/completions" elif (os.environ.get("DEEPSEEK_API_KEY") and resolved_key == os.environ.get("DEEPSEEK_API_KEY")) or (raw_model and "deepseek-chat" in raw_model): resolved_url = "https://api.deepseek.com/chat/completions" else: resolved_url = "https://openrouter.ai/api/v1/chat/completions" else: resolved_url = raw_url # Normalize URL: support base URLs (e.g. https://api.deepseek.com or https://api.deepseek.com/v1) and full endpoint resolved_url = resolved_url.rstrip("/") if not resolved_url.endswith("/chat/completions"): resolved_url = f"{resolved_url}/chat/completions" is_openrouter = "openrouter" in resolved_url.lower() or resolved_key.startswith("sk-or-") is_deepseek_official = "api.deepseek.com" in resolved_url.lower() # Model resolution resolved_model = raw_model if not resolved_model: resolved_model = "deepseek-chat" if is_deepseek_official else PE_MODEL_DEFAULT # If user selected DeepSeek official API, but left default OpenRouter model name (deepseek/deepseek-v4.1-flash), auto-switch to deepseek-chat if is_deepseek_official and ("deepseek/" in resolved_model.lower() or "flash" in resolved_model.lower()): resolved_model = "deepseek-chat" # Multimodal image packaging: DeepSeek official API only supports text; OpenRouter/OpenAI support vision content = [] if images and not is_deepseek_official: for image in images: # The enhancer receives JPEG: composite only this copy over white. # Original RGBA references keep their alpha for the diffusion VAE. rgba = ImageOps.exif_transpose(image).convert("RGBA") img_rgb = Image.new("RGB", rgba.size, "white") img_rgb.paste(rgba, mask=rgba.getchannel("A")) scale = min(1.0, (PE_MAX_PIXELS / (img_rgb.width * img_rgb.height)) ** 0.5) buf = io.BytesIO() img_rgb.resize( (max(32, int(img_rgb.width * scale)), max(32, int(img_rgb.height * scale))), Image.LANCZOS, ).save(buf, "JPEG", quality=90) content.append({ "type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()}, }) content.append({"type": "text", "text": user_text}) user_message_content = content elif images and is_deepseek_official: user_message_content = f"{user_text}\n(Note: User provided {len(images)} reference image(s) for visual edit context)." else: user_message_content = user_text body = { "model": resolved_model, "messages": [ {"role": "system", "content": PE_I2I if images else PE_T2I}, {"role": "user", "content": user_message_content}, ], } if is_openrouter: body["provider"] = {"sort": "throughput", "data_collection": "deny", "zdr": True} if not images: body["reasoning"] = {"effort": "low"} else: body["reasoning"] = {"enabled": False} headers = { "Authorization": f"Bearer {resolved_key}", "Content-Type": "application/json", } if is_openrouter: headers["HTTP-Referer"] = "https://huggingface.co/spaces" headers["X-Title"] = "Qwen Image 2.1 LoRA Studio" print(f"[pe] Requesting enhancement via {resolved_url} (model: {resolved_model})...", flush=True) try: with httpx.Client(timeout=30) as client: r = client.post(resolved_url, json=body, headers=headers) r.raise_for_status() data = r.json() if "choices" in data and len(data["choices"]) > 0: ans = data["choices"][0]["message"]["content"].strip() return ans or None, None return None, f"Empty choices response from {resolved_model}" except httpx.HTTPStatusError as e: err_body = "" try: err_body = f" - {e.response.text[:200]}" except Exception: pass err_msg = f"HTTP {e.response.status_code}{err_body}" print(f"[pe] failed: {err_msg}", flush=True) return None, err_msg except Exception as e: err_msg = f"{type(e).__name__}: {e}" print(f"[pe] failed: {err_msg}", flush=True) return None, err_msg def extract_pil_images(gallery_or_files): """Normalize reference uploads without resizing before the pipeline encodes them.""" if not gallery_or_files: return [] imgs = [] for item in gallery_or_files: img_obj = None if isinstance(item, (tuple, list)) and len(item) > 0: img_obj = item[0] elif isinstance(item, dict) and "image" in item: img_obj = item["image"] elif isinstance(item, (Image.Image, str)): img_obj = item if isinstance(img_obj, Image.Image): imgs.append(img_obj) elif isinstance(img_obj, str) and os.path.exists(img_obj): try: with Image.open(img_obj) as source: imgs.append(ImageOps.exif_transpose(source).convert("RGBA")) except Exception as e: raise gr.Error(f"无法读取参考图片：{Path(img_obj).name}") from e return imgs def manual_enhance_prompt( prompt: str, mode: str, ref_1: Image.Image | None, ref_2: Image.Image | None, multi_refs: list | None, aspect_ratio: str, api_key: str, model_name: str, base_url: str, ) -> str: """Manually enhance prompt and return the enriched string into UI.""" prompt = (prompt or "").strip() if not prompt: gr.Warning("Please enter a prompt first.") return prompt refs = [] if mode == "Edit Image (1 Ref)" and ref_1 is not None: refs = [ref_1] elif mode == "Transform & Swap (2 Refs)": refs = [img for img in [ref_1, ref_2] if img is not None] elif mode == "Multi-Image Reference (1-6 Refs)": refs = extract_pil_images(multi_refs) from rendering_config import AUTO, CUSTOM if aspect_ratio == AUTO and refs: last_ref = ImageOps.exif_transpose(refs[-1]) ratio_name = f"{last_ref.width}:{last_ref.height}" elif aspect_ratio == CUSTOM: ratio_name = "follow the user's composition; output dimensions are selected separately" else: ratio_name = aspect_ratio.split(" · ")[0] gr.Info("Enhancing prompt...") enhanced, pe_err = enhance_prompt( prompt=prompt, images=refs, ratio_name=ratio_name, api_key=api_key, model_name=model_name, base_url=base_url, ) if enhanced: gr.Info("✨ Prompt enhanced successfully!") return enhanced else: if pe_err: gr.Warning(f"Prompt enhancement failed: {pe_err}") else: gr.Warning("⚠️ Prompt enhancement returned empty result.") return prompt # ============================================================ # Inference Function # ============================================================ @gpu def _run_diffusion_pipeline( effective_prompt: str, width: int, height: int, steps: int, actual_seed: int, call_kwargs: dict, lora_slots: list[tuple[str, float, str, str]], cfg_scale: float = 1.0, negative_prompt: str = "", refine_v03: bool = False, ): plan = sampling_plan(lora_slots, steps, cfg_scale, refine_v03) merged_adapters = {} # adapter_name -> (eff_weight, display_name) for lora_name, lora_weight, c_repo, c_file in lora_slots: if not lora_name or lora_name == NONE_LORA or float(lora_weight) == 0: continue adapter_name, base_strength = ensure_adapter_ready(lora_name, c_repo, c_file) if adapter_name: eff_weight = float(lora_weight) * base_strength disp_name = lora_name if lora_name != "Custom HuggingFace LoRA..." else f"Custom({c_file or c_repo})" if adapter_name in merged_adapters: cur_w, cur_name = merged_adapters[adapter_name] merged_adapters[adapter_name] = (cur_w + eff_weight, cur_name) else: merged_adapters[adapter_name] = (eff_weight, disp_name) if merged_adapters: active_adapter_names = list(merged_adapters.keys()) active_adapter_weights = [v[0] for v in merged_adapters.values()] pipe.set_adapters(active_adapter_names, adapter_weights=active_adapter_weights) pipe.enable_lora() active_lora_desc = " + ".join([f"{v[1]} (scale={round(v[0], 2)})" for v in merged_adapters.values()]) else: pipe.disable_lora() active_lora_desc = "None" pipe_kwargs = dict(call_kwargs) cfg_val = plan["cfg"] pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config( BASE_SCHEDULER_CONFIG, **({"shift_terminal": None} if plan["is_viggle"] else {}), ) reextract[0] = False pipe_kwargs.update( sigmas=plan["sigmas"], output_resolution=1024, use_kv_cache=True, callback_on_step_end=base_tail if plan["refine"] else None, ) if cfg_val > 1.0: pipe_kwargs["true_cfg_scale"] = cfg_val pipe_kwargs["negative_prompt"] = negative_prompt.strip() if (negative_prompt and negative_prompt.strip()) else "" else: pipe_kwargs["true_cfg_scale"] = 1.0 try: with torch.inference_mode(): result = pipe( prompt=effective_prompt, width=width, height=height, num_inference_steps=plan["steps"], generator=torch.Generator("cuda").manual_seed(actual_seed), **pipe_kwargs, ).images[0] finally: # Restore even after cancellation/OOM during the base tail. reextract[0] = False if merged_adapters: pipe.enable_lora() return result, active_lora_desc def generate( prompt: str, negative_prompt: str = "", mode: str = "Text to Image", ref_image_1: Image.Image | None = None, ref_image_2: Image.Image | None = None, multi_refs: list | None = None, # LoRA Slot 1 lora_1: str = VIGGLE_V03, lora_1_strength: float = 1.0, custom_repo_1: str = "", custom_file_1: str = "", # LoRA Slot 2 lora_2: str = NONE_LORA, lora_2_strength: float = 1.0, custom_repo_2: str = "", custom_file_2: str = "", # LoRA Slot 3 lora_3: str = NONE_LORA, lora_3_strength: float = 1.0, custom_repo_3: str = "", custom_file_3: str = "", # LoRA Slot 4 lora_4: str = NONE_LORA, lora_4_strength: float = 1.0, custom_repo_4: str = "", custom_file_4: str = "", aspect_ratio: str = T2I_CHOICES[0], enhance_mode: str = "Auto", cfg_scale: float = 1.0, steps: int = 6, seed: int = 42, randomize_seed: bool = True, pe_api_key: str = "", pe_model: str = PE_MODEL_DEFAULT, pe_base_url: str = "", refine_v03: bool = False, custom_width: int = 1024, custom_height: int = 1024, progress: gr.Progress = gr.Progress(track_tqdm=True), ) -> tuple[str, str, int, str, str]: prompt = (prompt or "").strip() if not prompt: raise gr.Error("Please enter a prompt describing your image.") if len(prompt) > 4000: raise gr.Error("Prompt is too long. Please keep under 4000 characters.") if mode not in MODES: raise gr.Error(f"Invalid mode: {mode}") actual_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) if not 0 <= actual_seed <= MAX_SEED: raise gr.Error(f"Seed must be between 0 and {MAX_SEED}.") lora_slots = [ (lora_1, lora_1_strength, custom_repo_1, custom_file_1), (lora_2, lora_2_strength, custom_repo_2, custom_file_2), (lora_3, lora_3_strength, custom_repo_3, custom_file_3), (lora_4, lora_4_strength, custom_repo_4, custom_file_4), ] try: for name, strength, repo, filename in lora_slots: if name not in LORA_CHOICES: raise ValueError(f"Unknown LoRA: {name}") if name == "Custom HuggingFace LoRA..." and float(strength) > 0 and not (repo.strip() and filename.strip()): raise ValueError("请填写自定义 LoRA 的仓库和文件名。") plan = sampling_plan(lora_slots, steps, cfg_scale, refine_v03) except (ValueError, TypeError) as exc: raise gr.Error(str(exc)) from exc cfg_val = plan["cfg"] call_kwargs = {} refs_for_pe = [] if mode == "Edit Image (1 Ref)": if ref_image_1 is None: raise gr.Error("Please upload a reference image for single-image edit mode.") img = ImageOps.exif_transpose(ref_image_1).convert("RGBA") call_kwargs["image"] = img refs_for_pe = [ref_image_1] elif mode == "Transform & Swap (2 Refs)": if ref_image_1 is None or ref_image_2 is None: raise gr.Error("Please upload both Image 1 (Base) and Image 2 (Donor/Pose) for this mode.") img1 = ImageOps.exif_transpose(ref_image_1).convert("RGBA") img2 = ImageOps.exif_transpose(ref_image_2).convert("RGBA") call_kwargs["image"] = [img1, img2] refs_for_pe = [ref_image_1, ref_image_2] elif mode == "Multi-Image Reference (1-6 Refs)": multi_imgs = extract_pil_images(multi_refs) if not multi_imgs: raise gr.Error("Please upload at least one reference image for Multi-Image Reference mode.") if len(multi_imgs) > 6: raise gr.Error("最多支持 6 张参考图，请移除多余图片。") processed_imgs = [] for raw_img in multi_imgs: m_img = ImageOps.exif_transpose(raw_img).convert("RGBA") processed_imgs.append(m_img) call_kwargs["image"] = processed_imgs if len(processed_imgs) > 1 else processed_imgs[0] refs_for_pe = multi_imgs processed_refs = call_kwargs.get("image", []) if isinstance(processed_refs, Image.Image): processed_refs = [processed_refs] refs_for_pe = processed_refs for name, strength, _, _ in lora_slots: if float(strength) > 0 and ADAPTER_SPECS.get(name, {}).get("requires_two_images") and len(processed_refs) != 2: raise gr.Error(f"{name} 需要恰好两张参考图，请切换双图模式或上传两张参考图。") try: width, height, size_note = resolve_size( aspect_ratio, bool(processed_refs), custom_width, custom_height, processed_refs[-1].size if processed_refs else None, ) except (ValueError, TypeError) as exc: raise gr.Error(str(exc)) from exc ratio_name = f"{width}:{height}" # Step 1: Prompt Enhancement (runs before GPU acquisition to conserve ZeroGPU quota) used_prompt = prompt enhance_note = "" pe_time = 0.0 should_enhance = False if prompt.strip(): if enhance_mode == "On": should_enhance = True elif enhance_mode == "Auto" and n_tokens(prompt) < PE_AUTO_TOKENS: should_enhance = True if should_enhance: pe_start = time.perf_counter() enhanced, pe_err = enhance_prompt( prompt=prompt, images=refs_for_pe, ratio_name=ratio_name, api_key=pe_api_key, model_name=pe_model, base_url=pe_base_url, ) pe_time = time.perf_counter() - pe_start if enhanced and enhanced != prompt: used_prompt = enhanced enhance_note = f" | 🪄 Enhanced ({pe_time:.1f}s)" elif pe_err: enhance_note = f" | ⚠️ PE failed: {pe_err[:40]}" else: enhance_note = " | 🪄 Enhancement skipped/as written" effective_prompt = used_prompt # Step 2: Denoise on GPU start_time = time.perf_counter() result, active_lora_desc = _run_diffusion_pipeline( effective_prompt=effective_prompt, width=width, height=height, steps=int(steps), actual_seed=actual_seed, call_kwargs=call_kwargs, lora_slots=lora_slots, cfg_scale=cfg_val, negative_prompt=negative_prompt, refine_v03=refine_v03, ) elapsed = time.perf_counter() - start_time # Embed metadata into PNG chunks metadata = { "model": MODEL_ID, "precision": "bfloat16", "model_revision": MODEL_REVISION, "lora_adapter": active_lora_desc, "original_prompt": prompt, "prompt": effective_prompt, "negative_prompt": negative_prompt if cfg_val > 1.0 else "", "mode": mode, "seed": actual_seed, "steps": plan["steps"], "raw_sigmas": plan["sigmas"], "refine_v03": plan["refine"], "cfg_scale": cfg_val, "dimensions": f"{result.width}x{result.height}", "elapsed_seconds": round(elapsed, 2), "pe_seconds": round(pe_time, 2) if pe_time > 0 else 0, } png_info = PngImagePlugin.PngInfo() png_info.add_text("parameters", json.dumps(metadata, ensure_ascii=False)) temp_dir = tempfile.mkdtemp(prefix="qwen_aio_lora_") out_png_path = os.path.join(temp_dir, f"qwen_aio_{actual_seed}.png") result.save(out_png_path, "PNG", pnginfo=png_info, optimize=True) cfg_info = f" | CFG: {cfg_val}" if cfg_val > 1.0 else " | CFG: 1.0 (Native)" details = ( f"⚡ Time: {elapsed:.2f}s{enhance_note} | Seed: {actual_seed} | Steps: {plan['steps']}{cfg_info}\n" f"🎛️ LoRA: {active_lora_desc} | Size: {result.width}x{result.height}\n" f"{'7 Turbo + 2 Base' if plan['refine'] else 'Standard'} {size_note}" ) return out_png_path, out_png_path, actual_seed, details, effective_prompt # ============================================================ # Interface # ============================================================ from studio_ui import build_ui, CUSTOM_CSS demo = build_ui(globals()) if __name__ == "__main__": demo.queue(max_size=16, default_concurrency_limit=1).launch( mcp_server=True, css=CUSTOM_CSS, )