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) else: gpu = lambda fn: fn import httpx import gradio as gr import torch from accelerate import init_empty_weights from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel from diffusers.models.model_loading_utils import load_gguf_checkpoint from diffusers.quantizers.gguf.utils import dequantize_gguf_tensor from huggingface_hub import hf_hub_download from PIL import Image, ImageOps, PngImagePlugin from safetensors.torch import load_file as safetensors_load_file # Model Configuration MODEL_ID = "KasugaiSakura/Qwen-Image-2.1-Uncensored-Abenzerps-GGUF" CHECKPOINT = os.environ.get("QWEN_GGUF_CHECKPOINT", "qwen-image-2.1-UC-Q4_K_M.gguf") COMPANION_ID = "Qwen/Qwen-Image-2.1" COMPANION_REVISION = "790c92633540aa0cb11d9abf19eb46d861714758" SHA256_CHECKSUMS = { # Uncensored (UC) GGUFs "qwen-image-2.1-UC-Q4_K_M.gguf": "e79c8a009f2ecbdb6c70fd663d9aea9ee304a0d91f347e4169a756b8ad141b41", "qwen-image-2.1-UC-Q4_0.gguf": "13f59f20656efc0aa385d03c1fcac1a9dc2ad6e5ccc0ea9bfe1d6ac636f2c5b9", "qwen-image-2.1-UC-Q5_K_M.gguf": "af0bf278cf16d204fb31c384dc82fd41dca82d976b15fe9305a60c726fd6f821", "qwen-image-2.1-UC-Q6_K.gguf": "e14bb312109333b3d73b92ad9b9ac8b29b51f2edca1e7b86d1af1981bf1c4ee3", "qwen-image-2.1-UC-Q8_0.gguf": "cde456c72ea3ecebfc1be783300e972711d875e0c5f1bed33d42b66b156affa8", "qwen-image-2.1-UC-BF16.gguf": "f151c683a8aed4b310777017ebbbe3f2180f1180f7867115171adb7d50b0762a", } MODES = [ "Text to Image", "Multi-Image Reference (1-6 Refs)", "Edit Image (1 Ref)", "Transform & Swap (2 Refs)", "Transparent PNG", ] SIZES = { "Square · 1:1 (1024x1024)": (1024, 1024), "Landscape · 16:9 (1344x768)": (1344, 768), "Portrait · 9:16 (768x1344)": (768, 1344), "Landscape · 4:3 (1152x864)": (1152, 864), "Portrait · 3:4 (864x1152)": (864, 1152), } MAX_SEED = 2**31 - 1 # ============================================================ # Curated All-In-One LoRA Specifications # ============================================================ NONE_LORA = "None (Base Uncensored)" ADAPTER_SPECS = { "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() # ============================================================ # GGUF Checkpoint Initialization on CUDA # ============================================================ print(f"Loading checkpoint {CHECKPOINT} from {MODEL_ID}...", flush=True) checkpoint_path = hf_hub_download(MODEL_ID, CHECKPOINT) if CHECKPOINT in SHA256_CHECKSUMS: with open(checkpoint_path, "rb") as f: checksum = hashlib.file_digest(f, "sha256").hexdigest() if checksum != SHA256_CHECKSUMS[CHECKPOINT]: raise RuntimeError(f"Checksum verification failed for {CHECKPOINT}! Got {checksum}") print(f"Checksum verified: {checksum}", flush=True) # Expand the GGUF weights to bfloat16 once at startup for ZeroGPU eager packing weights = load_gguf_checkpoint(checkpoint_path) for name in list(weights.keys()): weights[name] = dequantize_gguf_tensor(weights[name]).to(torch.bfloat16) config = QwenImage21Transformer2DModel.load_config( COMPANION_ID, subfolder="transformer", revision=COMPANION_REVISION ) with init_empty_weights(): transformer = QwenImage21Transformer2DModel.from_config(config) transformer.load_state_dict(weights, strict=True, assign=True) transformer.eval().requires_grad_(False) print(f"Loaded {len(weights)} tensors into QwenImage21Transformer2DModel.", flush=True) del weights # Load full pipeline with companion VAE and text encoder, eagerly placed on 'cuda' pipe = QwenImage21Pipeline.from_pretrained( COMPANION_ID, revision=COMPANION_REVISION, transformer=transformer, torch_dtype=torch.bfloat16, ).to("cuda") print(f"{MODEL_ID} successfully initialized on CUDA for ZeroGPU.", flush=True) # ============================================================ # 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, float(spec.get("default_strength", 1.0)) # ============================================================ # 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() or os.environ.get("HF_TOKEN", "").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: img_rgb = ImageOps.exif_transpose(image).convert("RGB") 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 Uncensored 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 Gradio gallery/file inputs into a clean list of RGBA PIL Images.""" 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: imgs.append(ImageOps.exif_transpose(Image.open(img_obj)).convert("RGBA")) except Exception as e: print(f"[multi_ref] Failed to read image {img_obj}: {e}", flush=True) 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) ratio_name = aspect_ratio.split(" · ")[1].split(" ")[0] if " · " in aspect_ratio else "1:1" 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 = "", ): 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: 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) 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 = float(cfg_scale) if cfg_scale is not None else 1.0 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 with torch.inference_mode(): result = pipe( prompt=effective_prompt, width=width, height=height, num_inference_steps=int(steps), generator=torch.Generator("cuda").manual_seed(actual_seed), **pipe_kwargs, ).images[0] 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 = NONE_LORA, 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 = "Square · 1:1 (1024x1024)", enhance_mode: str = "Auto", cfg_scale: float = 1.0, steps: int = 30, seed: int = 42, randomize_seed: bool = True, pe_api_key: str = "", pe_model: str = PE_MODEL_DEFAULT, pe_base_url: str = "", 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}") if aspect_ratio not in SIZES: raise gr.Error(f"Invalid aspect ratio: {aspect_ratio}") if steps is None or not (4 <= int(steps) <= 40): raise gr.Error("Inference steps must be between 4 and 40.") actual_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) width, height = SIZES[aspect_ratio] ratio_name = aspect_ratio.split(" · ")[1].split(" ")[0] if " · " in aspect_ratio else "1:1" cfg_val = float(cfg_scale) if cfg_scale is not None else 1.0 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") img.thumbnail((2048, 2048)) 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") img1.thumbnail((2048, 2048)) img2 = ImageOps.exif_transpose(ref_image_2).convert("RGBA") img2.thumbnail((2048, 2048)) 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: gr.Warning(f"Only the first 6 reference images will be used (received {len(multi_imgs)}).") multi_imgs = multi_imgs[:6] processed_imgs = [] for raw_img in multi_imgs: m_img = ImageOps.exif_transpose(raw_img).convert("RGBA") m_img.thumbnail((2048, 2048)) processed_imgs.append(m_img) call_kwargs["image"] = processed_imgs if len(processed_imgs) > 1 else processed_imgs[0] refs_for_pe = multi_imgs # 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 if mode == "Transparent PNG": effective_prompt = ( "This is an RGBA image with transparency. " + used_prompt + " The image has alpha channel and the background is transparent." ) 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), ] # 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, ) elapsed = time.perf_counter() - start_time # Embed metadata into PNG chunks metadata = { "model": MODEL_ID, "checkpoint": CHECKPOINT, "lora_adapter": active_lora_desc, "original_prompt": prompt, "prompt": used_prompt, "negative_prompt": negative_prompt if cfg_val > 1.0 else "", "mode": mode, "seed": actual_seed, "steps": int(steps), "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: {steps}{cfg_info}\n" f"🎛️ LoRA: {active_lora_desc} | Size: {result.width}x{result.height}\n" f"🔒 Privacy: Zero data retention (session ephemeral)" ) return out_png_path, out_png_path, actual_seed, details, used_prompt # ============================================================ # UI Helpers # ============================================================ def update_lora_selection(selected_lora: str, current_prompt: str, current_steps: int): spec = ADAPTER_SPECS.get(selected_lora) show_custom = gr.update(visible=(selected_lora == "Custom HuggingFace LoRA...")) if not spec: return current_prompt, current_steps, 1.0, show_custom preset = spec.get("preset_prompt", "") new_prompt = current_prompt if preset and not current_prompt.strip(): new_prompt = preset elif preset and preset not in current_prompt: new_prompt = f"{current_prompt}, {preset}" if current_prompt.strip() else preset recommended_steps = spec.get("default_steps", current_steps) recommended_strength = spec.get("default_strength", 1.0) return new_prompt, recommended_steps, recommended_strength, show_custom def update_mode_ui(mode: str): is_multi = mode == "Multi-Image Reference (1-6 Refs)" is_edit_1 = mode == "Edit Image (1 Ref)" is_swap_2 = mode == "Transform & Swap (2 Refs)" return ( gr.update(visible=is_multi), gr.update(visible=(is_edit_1 or is_swap_2)), gr.update(visible=is_swap_2), ) # ============================================================ # Gradio Interface # ============================================================ CUSTOM_CSS = """ .gradio-container { max-width: 1280px !important; margin: 0 auto !important; } .badge { display: inline-block; padding: 2px 8px; border-radius: 6px; font-size: 0.8rem; font-weight: 600; margin-right: 6px; } .badge-turbo { background: #fee2e2; color: #991b1b; } .badge-style { background: #ede9fe; color: #5b21b6; } .badge-tool { background: #e0f2fe; color: #075985; } #generate-btn { font-weight: 700; font-size: 1.1rem; } """ with gr.Blocks(title="Qwen Image 2.1 Uncensored All-In-One LoRA Studio", css=CUSTOM_CSS, delete_cache=(3600, 86400)) as demo: gr.Markdown( "# 🚀 Qwen Image 2.1 Uncensored All-In-One LoRA Studio\n" "### Supercharged Text-to-Image, Image Editing & Guided Synthesis powered by ZeroGPU\n" "Uncensored Base (`KasugaiSakura/Qwen-Image-2.1-Uncensored-Abenzerps-GGUF`) + 15+ On-Demand Style, Speed, and Face/Pose LoRAs" ) with gr.Row(): with gr.Column(scale=6): mode_selector = gr.Radio( choices=MODES, value=MODES[0], label="Generation Mode", ) prompt_input = gr.Textbox( label="Prompt", placeholder="Describe your vision, character, scene, or requested edit...", lines=3, max_lines=6, ) with gr.Row(): enhance_radio = gr.Radio( choices=["Auto", "On", "Off"], value="Auto", label="Enhance Prompt (Viggle Turbo / DeepSeek V4.1 Flash)", info=f"Rewrites the prompt using DeepSeek V4.1 Flash (+~2s). 'Auto' rewrites prompts under {PE_AUTO_TOKENS} tokens and sends longer ones as written; 'Off' sends nothing.", scale=4, ) manual_enhance_btn = gr.Button("🪄 Enhance Prompt", size="sm", scale=1) with gr.Row(): multi_refs_gallery = gr.Gallery( label="Reference Images (Upload 1-6 images, refer to as Image 1, Image 2... in prompt)", type="pil", columns=3, height=280, interactive=True, visible=False, ) with gr.Row(): ref_img_1 = gr.Image( label="Image 1 (Base / Target)", type="pil", visible=False, ) ref_img_2 = gr.Image( label="Image 2 (Face Donor / Pose Reference)", type="pil", visible=False, ) with gr.Accordion("🎨 Multi-LoRA Studio (Stack up to 4 LoRAs)", open=True): gr.Markdown("💡 **Multi-LoRA Stacking**: Combine multiple adapters simultaneously (e.g. Turbo Speed + Art Style + Face Swap + Detail Enhancer).") # Slot 1 with gr.Group(): with gr.Row(): lora_1_dropdown = gr.Dropdown( choices=LORA_CHOICES, value=NONE_LORA, label="LoRA Slot 1 (Primary)", info="Select your primary style, speed booster, or transfer LoRA", scale=3, ) lora_1_strength = gr.Slider( minimum=0.0, maximum=1.5, value=1.0, step=0.05, label="Weight 1", scale=1, ) with gr.Row(visible=False) as custom_box_1: custom_repo_1 = gr.Textbox(label="LoRA 1 HF Repo", placeholder="e.g. baka999/Q2.1") custom_file_1 = gr.Textbox(label="LoRA 1 Weights File", placeholder="e.g. NSFW Qwen by TheseAlpacas V2.safetensors") # Additional Slots 2, 3, 4 with gr.Accordion("➕ Stack Additional LoRAs (Slot 2, 3, 4)", open=False): with gr.Group(): with gr.Row(): lora_2_dropdown = gr.Dropdown( choices=LORA_CHOICES, value=NONE_LORA, label="LoRA Slot 2", scale=3, ) lora_2_strength = gr.Slider( minimum=0.0, maximum=1.5, value=1.0, step=0.05, label="Weight 2", scale=1, ) with gr.Row(visible=False) as custom_box_2: custom_repo_2 = gr.Textbox(label="LoRA 2 HF Repo", placeholder="e.g. reverentelusarca/elusarcas-qwen-2.1-detail-enhancer-lora") custom_file_2 = gr.Textbox(label="LoRA 2 Weights File", placeholder="e.g. elusarcas-qwen2-1-detailer-v1.safetensors") with gr.Group(): with gr.Row(): lora_3_dropdown = gr.Dropdown( choices=LORA_CHOICES, value=NONE_LORA, label="LoRA Slot 3", scale=3, ) lora_3_strength = gr.Slider( minimum=0.0, maximum=1.5, value=1.0, step=0.05, label="Weight 3", scale=1, ) with gr.Row(visible=False) as custom_box_3: custom_repo_3 = gr.Textbox(label="LoRA 3 HF Repo", placeholder="e.g. WarmBloodAban/Qwen-Image-2.1-LoRAs") custom_file_3 = gr.Textbox(label="LoRA 3 Weights File", placeholder="e.g. Qwen2.1_Anything2RealCharacters.safetensors") with gr.Group(): with gr.Row(): lora_4_dropdown = gr.Dropdown( choices=LORA_CHOICES, value=NONE_LORA, label="LoRA Slot 4", scale=3, ) lora_4_strength = gr.Slider( minimum=0.0, maximum=1.5, value=1.0, step=0.05, label="Weight 4", scale=1, ) with gr.Row(visible=False) as custom_box_4: custom_repo_4 = gr.Textbox(label="LoRA 4 HF Repo", placeholder="e.g. ausboss/Qwen-Image-2.1-Outpaint-LoRA") custom_file_4 = gr.Textbox(label="LoRA 4 Weights File", placeholder="e.g. qwen-image-2.1-outpaint-v2.safetensors") with gr.Accordion("⚙️ Advanced Generation Settings", open=False): with gr.Row(): aspect_ratio_dropdown = gr.Dropdown( choices=list(SIZES.keys()), value=list(SIZES.keys())[0], label="Aspect Ratio", ) steps_slider = gr.Slider( minimum=4, maximum=40, value=30, step=1, label="Inference Steps", info="4-6 steps for Turbo LoRAs, 25-35 for regular generation", ) with gr.Row(): cfg_slider = gr.Slider( minimum=1.0, maximum=10.0, value=1.0, step=0.1, label="CFG Scale (true_cfg_scale)", info="1.0 = native guidance-free (fastest, recommended for Turbo/Qwen 2.1). > 1.0 activates negative prompt guidance.", ) seed_number = gr.Number(value=42, label="Seed", precision=0) randomize_seed_cb = gr.Checkbox(value=True, label="Randomize Seed") negative_prompt_input = gr.Textbox( label="Negative Prompt (Active when CFG > 1.0)", placeholder="e.g. blurry, low quality, distorted, bad anatomy, deformed limbs...", lines=2, max_lines=4, ) with gr.Accordion("🤖 Prompt Enhancer API Settings (Optional)", open=False): pe_api_key_input = gr.Textbox( label="OpenRouter / OpenAI / DeepSeek API Key", type="password", placeholder="Optional if OPENROUTER_API_KEY, OPENAI_API_KEY, or DEEPSEEK_API_KEY is in env", ) with gr.Row(): pe_model_input = gr.Textbox( label="Enhancer Model", value=PE_MODEL_DEFAULT, placeholder="deepseek/deepseek-v4.1-flash (OpenRouter) or deepseek-chat (DeepSeek API)", ) pe_base_url_input = gr.Textbox( label="Custom API Base URL", placeholder="e.g. https://api.deepseek.com/v1 (or https://openrouter.ai/api/v1)", ) generate_btn = gr.Button( "✨ Generate Image", variant="primary", elem_id="generate-btn", ) with gr.Column(scale=6): output_image = gr.Image( label="Generated Output", type="filepath", interactive=False, ) with gr.Row(): download_file = gr.File( label="Download High-Res PNG (Includes Parameters)", interactive=False, ) used_prompt_box = gr.Textbox( label="Prompt Sent to the Model", lines=3, interactive=False, info="The final prompt sent to the image model (shows enriched prompt when enhanced).", ) details_box = gr.Textbox( label="Execution Details & Privacy", interactive=False, ) # Event Bindings mode_selector.change( fn=update_mode_ui, inputs=[mode_selector], outputs=[multi_refs_gallery, ref_img_1, ref_img_2], ) lora_1_dropdown.change( fn=update_lora_selection, inputs=[lora_1_dropdown, prompt_input, steps_slider], outputs=[prompt_input, steps_slider, lora_1_strength, custom_box_1], ) lora_2_dropdown.change( fn=update_lora_selection, inputs=[lora_2_dropdown, prompt_input, steps_slider], outputs=[prompt_input, steps_slider, lora_2_strength, custom_box_2], ) lora_3_dropdown.change( fn=update_lora_selection, inputs=[lora_3_dropdown, prompt_input, steps_slider], outputs=[prompt_input, steps_slider, lora_3_strength, custom_box_3], ) lora_4_dropdown.change( fn=update_lora_selection, inputs=[lora_4_dropdown, prompt_input, steps_slider], outputs=[prompt_input, steps_slider, lora_4_strength, custom_box_4], ) manual_enhance_btn.click( fn=manual_enhance_prompt, inputs=[ prompt_input, mode_selector, ref_img_1, ref_img_2, multi_refs_gallery, aspect_ratio_dropdown, pe_api_key_input, pe_model_input, pe_base_url_input, ], outputs=[prompt_input], ) generate_btn.click( fn=generate, inputs=[ prompt_input, negative_prompt_input, mode_selector, ref_img_1, ref_img_2, multi_refs_gallery, lora_1_dropdown, lora_1_strength, custom_repo_1, custom_file_1, lora_2_dropdown, lora_2_strength, custom_repo_2, custom_file_2, lora_3_dropdown, lora_3_strength, custom_repo_3, custom_file_3, lora_4_dropdown, lora_4_strength, custom_repo_4, custom_file_4, aspect_ratio_dropdown, enhance_radio, cfg_slider, steps_slider, seed_number, randomize_seed_cb, pe_api_key_input, pe_model_input, pe_base_url_input, ], outputs=[output_image, download_file, seed_number, details_box, used_prompt_box], api_name="generate", concurrency_limit=1, concurrency_id="qwen-aio-pipeline", ) gr.Examples( examples=[ [ "A futuristic neon cyberpunk samurai in a rain-soaked Tokyo alleyway, glowing katana, volumetric mist", "Text to Image", "Anime Consistency", ], [ "A tranquil highland mountain lake surrounded by autumn pines at golden hour, reflections in water", "Text to Image", "Natural Exposure (Photorealism)", ], [ "A charming little steampunk robot holding a delicate glass flower, highly detailed gears and brass clockwork", "Text to Image", "Viggle Turbo (4-Step Acceleration)", ], ], inputs=[prompt_input, mode_selector, lora_1_dropdown], label="Try an Idea", ) if __name__ == "__main__": demo.queue(max_size=16, default_concurrency_limit=1).launch( mcp_server=True, )