Qdpie's picture
Upload 74 files
09472e9 verified
Raw
History Blame Contribute Delete
46.7 kB
import os
import gc
import gradio as gr
import numpy as np
import spaces
import torch
import random
import base64
import json
import html as html_lib
from io import BytesIO
from PIL import Image
MAX_SEED = np.iinfo(np.int32).max
LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("torch.__version__ =", torch.__version__)
print("torch.version.cuda =", torch.version.cuda)
print("cuda available:", torch.cuda.is_available())
print("cuda device count:", torch.cuda.device_count())
if torch.cuda.is_available():
print("current device:", torch.cuda.current_device())
print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
print("Using device:", device)
from diffusers import FlowMatchEulerDiscreteScheduler
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
dtype = torch.bfloat16
# ---------------------------------------------------------------------------
# Model load strategy (keeps cold-start near the ~20GB Rapid path, not ~58GB):
# 1) Load Rapid-AIO-V19 transformer (~20GB FP8) β€” this is the diffusion brain.
# 2) Load Qwen-Image-Edit-2511 pipeline shell (text_encoder, VAE, processor…)
# while NEVER downloading the base BF16 transformer/* weights (~40GB).
# Passing transformer= alone is usually enough for diffusers to skip that
# component, but a factory rebuild / cache miss can still pull the whole
# Qwen/Qwen-Image-Edit-2511 repo (~57.7GB on the Hub). snapshot_download with
# ignore_patterns makes the skip explicit and log-friendly.
# ---------------------------------------------------------------------------
from huggingface_hub import snapshot_download
print("=== [1/2] Loading Rapid transformer (prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19, ~20GB) ===")
_transformer = QwenImageTransformer2DModel.from_pretrained(
"prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
torch_dtype=dtype,
device_map="cuda" if torch.cuda.is_available() else None,
)
print("=== [2/2] Loading Qwen-Image-Edit-2511 shell (skip base transformer/* weights) ===")
_base_dir = snapshot_download(
"Qwen/Qwen-Image-Edit-2511",
# Keep transformer/config.json if present for schema, but never the multi-GB weights.
ignore_patterns=[
"transformer/*.safetensors",
"transformer/*.bin",
"transformer/*.pt",
"transformer/*.pth",
"transformer/*.index.json",
"transformer/**/*.safetensors",
"transformer/**/*.bin",
],
)
import warnings
with warnings.catch_warnings():
# Local qwenimage.transformer_qwenimage is API-compatible with the official
# diffusers class; type-check warning is noisy but harmless.
warnings.filterwarnings(
"ignore",
message=r".*Expected types for transformer.*",
)
pipe = QwenImageEditPlusPipeline.from_pretrained(
_base_dir,
transformer=_transformer,
torch_dtype=dtype,
).to(device)
print("=== Pipeline ready (Rapid transformer + base text_encoder/VAE) ===")
try:
_attn_proc = QwenDoubleStreamAttnProcessorFA3()
pipe.transformer.set_attn_processor(_attn_proc)
_backend = getattr(_attn_proc, "_attention_backend", "unknown")
if _backend == "fa3":
print("Attention backend: Flash Attention 3")
else:
print(f"Attention backend: {_backend} (FA3 unavailable on this GPU β€” expected on Blackwell)")
except Exception as e:
print(f"Warning: Could not set attention processor: {e}")
ADAPTER_SPECS = {
"Multiple-Angles": {
"repo": "dx8152/Qwen-Edit-2509-Multiple-angles",
"weights": "ι•œε€΄θ½¬ζ’.safetensors",
"adapter_name": "multiple-angles",
},
"Fal-Multiple-Angles": {
"repo": "fal/Qwen-Image-Edit-2511-Multiple-Angles-LoRA",
"weights": "qwen-image-edit-2511-multiple-angles-lora.safetensors",
"adapter_name": "fal-multiple-angles",
},
"Photo-to-Anime": {
"repo": "autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime",
"weights": "Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors",
"adapter_name": "photo-to-anime",
},
"Anime-V2": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Anime",
"weights": "Qwen-Image-Edit-2511-Anime-2000.safetensors",
"adapter_name": "anime-v2",
},
"Manga-Tone": {
"repo": "nappa114514/Qwen-Image-Edit-2509-Manga-Tone",
"weights": "tone001.safetensors",
"adapter_name": "manga-tone",
},
"Noir-Comic-Book": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Noir-Comic-Book-Panel",
"weights": "Noir-Comic-Book-Panel_20.safetensors",
"adapter_name": "ncb",
},
"Pixar-Inspired-3D": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Pixar-Inspired-3D",
"weights": "PI3_20.safetensors",
"adapter_name": "pi3",
},
"Polaroid-Photo": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Polaroid-Photo",
"weights": "Qwen-Image-Edit-2511-Polaroid-Photo.safetensors",
"adapter_name": "polaroid-photo",
},
"Hyper-Realistic-Portrait": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Hyper-Realistic-Portrait",
"weights": "HRP_20.safetensors",
"adapter_name": "hyper-realistic-portrait",
},
"Ultra-Realistic-Portrait": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Ultra-Realistic-Portrait",
"weights": "URP_20.safetensors",
"adapter_name": "ultra-realistic-portrait",
},
"Anything2Real": {
"repo": "lrzjason/Anything2Real_2601",
"weights": "anything2real_2601.safetensors",
"adapter_name": "anything2real",
},
"Style-Transfer": {
"repo": "zooeyy/Style-Transfer",
"weights": "Style Transfer-Alpha-V0.1.safetensors",
"adapter_name": "style-transfer",
},
"Upscaler": {
"repo": "starsfriday/Qwen-Image-Edit-2511-Upscale2K",
"weights": "qwen_image_edit_2511_upscale.safetensors",
"adapter_name": "upscale-2k",
},
"Unblur-Anything": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Unblur-Upscale",
"weights": "Qwen-Image-Edit-Unblur-Upscale_15.safetensors",
"adapter_name": "unblur-anything",
},
"Light-Migration": {
"repo": "dx8152/Qwen-Edit-2509-Light-Migration",
"weights": "参考色调.safetensors",
"adapter_name": "light-migration",
},
"Any-light": {
"repo": "lilylilith/QIE-2511-MP-AnyLight",
"weights": "QIE-2511-AnyLight_.safetensors",
"adapter_name": "any-light",
},
"Studio-DeLight": {
"repo": "prithivMLmods/QIE-2511-Studio-DeLight",
"weights": "QIE-2511-Studio-DeLight-5000.safetensors",
"adapter_name": "studio-delight",
},
"Cinematic-FlatLog": {
"repo": "prithivMLmods/QIE-2511-Cinematic-FlatLog-Control",
"weights": "QIE-2511-Cinematic-FlatLog-Control-3200.safetensors",
"adapter_name": "flat-log",
},
"Midnight-Noir-Eyes-Spotlight": {
"repo": "prithivMLmods/Qwen-Image-Edit-2511-Midnight-Noir-Eyes-Spotlight",
"weights": "Qwen-Image-Edit-2511-Midnight-Noir-Eyes-Spotlight.safetensors",
"adapter_name": "midnight-noir-eyes-spotlight",
},
}
LOADED_ADAPTERS: set = set()
ADAPTER_NAMES = list(ADAPTER_SPECS.keys())
LORA_DISPLAY_ORDER = ["Multiple-Angles", "Fal-Multiple-Angles", "Photo-to-Anime", "Anime-V2", "Manga-Tone", "Noir-Comic-Book", "Pixar-Inspired-3D", "Polaroid-Photo", "Hyper-Realistic-Portrait", "Ultra-Realistic-Portrait", "Anything2Real", "Style-Transfer", "Upscaler", "Unblur-Anything", "Light-Migration", "Any-light", "Studio-DeLight", "Cinematic-FlatLog", "Midnight-Noir-Eyes-Spotlight"]
LORA_HELP = {"none": "A LoRA is a specialized skill pack for the editor (style, camera, lighting, upscale, etc.). Pick one that matches your goal, then write a prompt that asks for that kind of change. Use LoRA Prompts chips for starter text that matches each adapter. Prompt and LoRA work together β€” the menu only selects the skill pack.", "Multiple-Angles": "Camera / viewpoint LoRA. Best for rotating or shifting perspective (e.g. rotate 45 degrees right, top-down view). Use with Camera quick prompts.", "Fal-Multiple-Angles": "Another multi-angle camera LoRA (2511-oriented). Viewpoint changes driven by the prompt (e.g. front-right quarter view). Try if Multiple-Angles is weak.", "Photo-to-Anime": "Turns photos toward anime illustration. Pair with prompts like transform into anime.", "Anime-V2": "Alternate anime conversion LoRA with a different look than Photo-to-Anime.", "Manga-Tone": "Manga / comic tone rendering. Prompt e.g. paint with manga tone.", "Noir-Comic-Book": "Noir comic-book panel look. Prompt e.g. transform into a noir comic book style.", "Pixar-Inspired-3D": "Stylized 3D / Pixar-like rendering. Prompt e.g. transform into Pixar-inspired 3D.", "Polaroid-Photo": "Polaroid / instant photo aesthetic. Use polaroid snapshot language in the prompt.", "Hyper-Realistic-Portrait": "Hyper-real face / portrait detail. Best on faces.", "Ultra-Realistic-Portrait": "Realism portrait LoRA with a different look than Hyper-Realistic.", "Anything2Real": "Steers toward realistic photograph look. Prompt e.g. change to a realistic photograph.", "Style-Transfer": "Style from a reference image. Use 2 images: #1 content, #2 style.", "Upscaler": "Detail / resolution enhancement. Prompt e.g. upscale to 4K; larger output size helps.", "Unblur-Anything": "Sharpen / deblur soft images. Prompt e.g. unblur and upscale.", "Light-Migration": "Relight using a reference. Usually 2 images: subject + lighting reference.", "Any-light": "Lighting control. Often works with a lighting reference and a clear lighting prompt.", "Studio-DeLight": "Neutral / studio even lighting. Prompt e.g. neutral uniform lighting.", "Cinematic-FlatLog": "Cinematic flat / log color grade. Prompt e.g. cinematic flat log.", "Midnight-Noir-Eyes-Spotlight": "Dramatic noir spotlight look. Prompt for midnight noir eyes spotlight style."}
# Starter prompts when a LoRA is selected (creator-aligned). Edit freely.
# Applied on menu change if the prompt box is empty or still holds the previous default.
LORA_DEFAULT_PROMPTS = {
"none": "",
"Multiple-Angles": "Rotate the camera 45 degrees to the right.",
"Fal-Multiple-Angles": "Front-right quarter view.",
"Photo-to-Anime": "Transform into anime.",
"Anime-V2": "Transform into anime (while preserving the background and remaining elements maintaining realism and original details.)",
"Manga-Tone": "Paint with manga tone.",
"Noir-Comic-Book": "Transform into a noir comic book style.",
"Pixar-Inspired-3D": "Transform it into Pixar-inspired 3D.",
"Polaroid-Photo": "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed preserving realistic texture and details.",
"Hyper-Realistic-Portrait": "Transform into a hyper-realistic face portrait.",
"Ultra-Realistic-Portrait": "Ultra-realistic portrait.",
"Anything2Real": "Change the picture to realistic photograph.",
"Style-Transfer": "Convert Image 1 to the style of Image 2.",
"Upscaler": "Upscale this picture to 4K resolution.",
"Unblur-Anything": "Unblur and upscale.",
"Light-Migration": "Refer to the color tone, remove the original lighting from Image 1, and relight Image 1 based on the lighting and color tone of Image 2.",
"Any-light": "Apply the lighting from image 2 to image 1.",
"Studio-DeLight": "Neutral uniform lighting. Preserve identity and composition.",
"Cinematic-FlatLog": "Transform into a cinematic flat log.",
"Midnight-Noir-Eyes-Spotlight": "Transform into Midnight Noir Eyes Spotlight.",
}
# Demo thumbs live on disk under loras/ (not embedded) so app.py stays small / HF-editable.
# Gradio serves them via allowed_paths + /gradio_api/file=...
import json as _json
from pathlib import Path
_LORA_DIR = Path(__file__).resolve().parent / "loras"
_DEMO_EXTS = (".jpg", ".jpeg", ".png", ".webp")
def _gradio_file_url(path: Path) -> str:
"""URL Gradio uses for files under allowed_paths (absolute path is most reliable)."""
return f"/gradio_api/file={path.resolve()}"
def _find_demo(slug: str, kind: str):
"""kind is 'before' or 'after'. Returns Gradio file URL or None."""
for ext in _DEMO_EXTS:
candidate = _LORA_DIR / f"{slug}-{kind}{ext}"
if candidate.is_file():
return _gradio_file_url(candidate)
return None
def _build_demo_maps():
before, after = {}, {}
for name in LORA_DISPLAY_ORDER:
slug = name.lower()
b = _find_demo(slug, "before")
a = _find_demo(slug, "after")
if b:
before[slug] = b
if a:
after[slug] = a
gen_b = _find_demo("generic", "before") or ""
gen_a = _find_demo("generic", "after") or ""
return before, after, gen_b, gen_a
LORA_DEMO_BEFORE, LORA_DEMO_AFTER, LORA_DEMO_GENERIC_BEFORE, LORA_DEMO_GENERIC_AFTER = _build_demo_maps()
LORA_HELP_JS = _json.dumps(LORA_HELP, ensure_ascii=False)
LORA_DEFAULT_PROMPTS_JS = _json.dumps(LORA_DEFAULT_PROMPTS, ensure_ascii=False)
LORA_DEMO_BEFORE_JS = _json.dumps(LORA_DEMO_BEFORE)
LORA_DEMO_AFTER_JS = _json.dumps(LORA_DEMO_AFTER)
LORA_DEMO_GENERIC_BEFORE_JS = _json.dumps(LORA_DEMO_GENERIC_BEFORE)
LORA_DEMO_GENERIC_AFTER_JS = _json.dumps(LORA_DEMO_GENERIC_AFTER)
print(
f"LoRA demos from disk: {len(LORA_DEMO_BEFORE)} before, {len(LORA_DEMO_AFTER)} after "
f"(dir={_LORA_DIR})"
)
# ── Prompt Packs (inference-time instruction + reference bundles) ─────────────
# Activated by tags/phrases/keywords from each pack's pack.json β€” not training.
import re as _re
_PROMPT_PACKS_DIR = Path(__file__).resolve().parent / "prompt_packs"
_PROMPT_PACKS: list = []
def _load_prompt_packs() -> list:
packs = []
if not _PROMPT_PACKS_DIR.is_dir():
print(f"Prompt packs: no directory at {_PROMPT_PACKS_DIR}")
return packs
for sub in sorted(_PROMPT_PACKS_DIR.iterdir()):
if not sub.is_dir() or sub.name.startswith("."):
continue
cfg_path = sub / "pack.json"
if not cfg_path.is_file():
continue
try:
with open(cfg_path, "r", encoding="utf-8") as f:
cfg = json.load(f)
except Exception as e:
print(f"Prompt packs: failed to read {cfg_path}: {e}")
continue
if cfg.get("enabled", True) is False:
print(f"Prompt packs: skipped disabled '{sub.name}'")
continue
pack = {
"id": str(cfg.get("id") or sub.name),
"name": str(cfg.get("name") or sub.name),
"description": str(cfg.get("description") or "").strip(),
"priority": int(cfg.get("priority", 0)),
"triggers": cfg.get("triggers") or [],
"prompt_prefix": str(cfg.get("prompt_prefix") or ""),
"prompt_suffix": str(cfg.get("prompt_suffix") or ""),
"ref_images": cfg.get("ref_images") or [],
"strip_triggers": bool(cfg.get("strip_triggers", True)),
"force_lora": cfg.get("force_lora"),
"override_user_ref": bool(cfg.get("override_user_ref", True)),
"dir": sub,
}
# Resolve ref paths that exist
refs = []
for rel in pack["ref_images"]:
p = sub / rel
if p.is_file():
refs.append(p)
else:
print(f"Prompt packs: missing ref {p} (pack {pack['id']})")
pack["ref_paths"] = refs
packs.append(pack)
print(
f"Prompt packs: loaded '{pack['id']}' "
f"({len(pack['triggers'])} triggers, {len(refs)} refs)"
)
packs.sort(key=lambda p: (-p["priority"], p["id"]))
return packs
def _trigger_matches(trigger: dict, prompt: str, prompt_lower: str) -> bool:
if not isinstance(trigger, dict):
return False
t = (trigger.get("type") or "phrase").lower().strip()
raw = trigger.get("value", "")
if t == "always":
return bool(raw) if not isinstance(raw, bool) else raw
if raw is None or raw == "":
return False
val = str(raw)
val_lower = val.lower()
if t == "tag":
return val_lower in prompt_lower
if t == "phrase":
return val_lower in prompt_lower
if t == "keyword":
# word-boundary-ish match
return _re.search(r"(?<!\w)" + _re.escape(val_lower) + r"(?!\w)", prompt_lower) is not None
if t == "regex":
try:
return _re.search(val, prompt, flags=_re.IGNORECASE) is not None
except _re.error as e:
print(f"Prompt packs: bad regex {val!r}: {e}")
return False
print(f"Prompt packs: unknown trigger type {t!r}")
return False
def _pack_matches(pack: dict, prompt: str) -> list:
"""Return list of triggers that matched."""
prompt_lower = prompt.lower()
hit = []
for tr in pack.get("triggers") or []:
if _trigger_matches(tr, prompt, prompt_lower):
hit.append(tr)
return hit
def _strip_triggers_from_prompt(prompt: str, packs: list) -> str:
out = prompt
for pack in packs:
if not pack.get("strip_triggers", True):
continue
for tr in pack.get("triggers") or []:
if not isinstance(tr, dict):
continue
t = (tr.get("type") or "").lower()
val = tr.get("value")
if val is None or val == "" or t in ("always", "regex"):
continue
# Case-insensitive remove of tag/phrase/keyword
pattern = _re.escape(str(val))
out = _re.sub(pattern, " ", out, flags=_re.IGNORECASE)
out = _re.sub(r"[ \t]{2,}", " ", out)
out = _re.sub(r"\n{3,}", "\n\n", out)
return out.strip()
def apply_prompt_packs(prompt: str, pil_images: list, lora_adapter: str, menu_pack_id: str = "none"):
"""
Apply prompt packs: wrap prompt, optional ref as image #2, optional force LoRA.
Menu selection (menu_pack_id) overrides prompt triggers when not "none".
Returns (prompt, pil_images, lora_adapter, applied_ids).
"""
if not prompt:
return prompt, pil_images, lora_adapter, []
if not _PROMPT_PACKS:
return prompt, pil_images, lora_adapter, []
matched = []
menu_id = (menu_pack_id or "none").strip()
if menu_id and menu_id.lower() not in ("none", ""):
pack = next((p for p in _PROMPT_PACKS if p["id"] == menu_id), None)
if pack:
matched = [(pack, [{"type": "menu", "value": menu_id}])]
print(f"--- Prompt pack from menu (overrides tags): {menu_id} ---")
else:
print(f"Prompt packs: menu id {menu_id!r} not found; falling back to prompt triggers")
for p in _PROMPT_PACKS:
hits = _pack_matches(p, prompt)
if hits:
matched.append((p, hits))
else:
for pack in _PROMPT_PACKS:
hits = _pack_matches(pack, prompt)
if hits:
matched.append((pack, hits))
if not matched:
return prompt, pil_images, lora_adapter, []
applied_ids = [p["id"] for p, _ in matched]
# Menu path already logged; trigger path needs a line
if not (menu_id and menu_id.lower() not in ("none", "")):
print(f"--- Prompt packs applied (triggers): {', '.join(applied_ids)} ---")
# Build prompt wrap (all matched packs, high priority first β€” already sorted)
prefix_parts = []
suffix_parts = []
for pack, _ in matched:
if pack["prompt_prefix"]:
prefix_parts.append(pack["prompt_prefix"].strip())
if pack["prompt_suffix"]:
suffix_parts.append(pack["prompt_suffix"].strip())
body = _strip_triggers_from_prompt(prompt, [p for p, _ in matched])
pieces = []
if prefix_parts:
pieces.append(" ".join(prefix_parts))
if body:
pieces.append(body)
if suffix_parts:
pieces.append(" ".join(suffix_parts))
new_prompt = " ".join(pieces).strip() or body or prompt
# Reference image: highest-priority pack that has refs (list already priority-sorted)
new_images = list(pil_images)
ref_pack = next((p for p, _ in matched if p.get("ref_paths")), None)
if ref_pack and new_images:
try:
ref_img = Image.open(ref_pack["ref_paths"][0]).convert("RGB")
override = ref_pack.get("override_user_ref", True)
if len(new_images) == 1:
new_images = [new_images[0], ref_img]
print(f"Prompt pack '{ref_pack['id']}': attached ref as image #2")
elif len(new_images) >= 2 and override:
new_images = [new_images[0], ref_img]
print(f"Prompt pack '{ref_pack['id']}': replaced image #2 with pack ref")
else:
print(
f"Prompt pack '{ref_pack['id']}': ref not attached "
f"(user has 2 images and override_user_ref=false)"
)
except Exception as e:
print(f"Prompt pack '{ref_pack['id']}': failed to load ref: {e}")
# force_lora: first matched pack that sets it (highest priority)
new_lora = lora_adapter
for pack, _ in matched:
fl = pack.get("force_lora")
if fl and str(fl).strip() and str(fl).strip().lower() not in ("none", "null"):
name = str(fl).strip()
if name in ADAPTER_SPECS:
new_lora = name
print(f"Prompt pack '{pack['id']}': force_lora -> {name}")
break
else:
print(f"Prompt pack '{pack['id']}': force_lora {name!r} not in ADAPTER_SPECS")
return new_prompt, new_images, new_lora, applied_ids
_PROMPT_PACKS = _load_prompt_packs()
print(f"Prompt packs ready: {len(_PROMPT_PACKS)} enabled pack(s)")
PROMPT_PACK_IDS = [p["id"] for p in _PROMPT_PACKS]
_PROMPT_PACK_NONE_DESC = (
"Prompt Packs add variety the base model struggles with in specific areas. "
"Choose a pack in the menu below, or trigger one in the prompt (e.g. [[pack]])."
)
def _pack_force_lora_name(p: dict) -> str:
"""Return ADAPTER_SPECS key for pack force_lora, or '' if none/invalid."""
fl = p.get("force_lora")
if not fl or str(fl).strip().lower() in ("none", "null", ""):
return ""
name = str(fl).strip()
return name if name in ADAPTER_SPECS else ""
def _pack_option_html(p: dict) -> str:
vid = html_lib.escape(p["id"], quote=True)
fl_name = _pack_force_lora_name(p)
label = p["name"] + (f" Β· {p['id']}" if p["name"] != p["id"] else "")
if fl_name:
# Visible cue in the menu that this pack forces a LoRA on select/run
label = f"{label} Β· LoRA: {fl_name}"
label = html_lib.escape(label)
desc = p.get("description") or (
f"{p['name']}: fixed instructions"
+ (" + reference image #2" if p.get("ref_paths") else "")
+ ". Applies on Run (menu overrides prompt tags)."
)
if fl_name:
desc = f"{desc} Forces LoRA β€œ{fl_name}” (menu switches on select; wins on Run)."
desc_attr = html_lib.escape(desc, quote=True)
fl_attr = html_lib.escape(fl_name, quote=True)
return (
f'<option value="{vid}" data-description="{desc_attr}" '
f'data-force-lora="{fl_attr}">{label}</option>'
)
PROMPT_PACK_OPTIONS_HTML = "\n".join(
[
'<option value="none" selected data-description="{d}" data-force-lora="">None (use prompt tags if any)</option>'.format(
d=html_lib.escape(_PROMPT_PACK_NONE_DESC, quote=True)
)
]
+ [_pack_option_html(p) for p in _PROMPT_PACKS]
)
EXAMPLES_CONFIG = [
{"images": ["examples/B.jpg"], "prompt": "Transform into anime.", "lora": "Photo-to-Anime"},
{"images": ["examples/HRP.jpg"], "prompt": "Transform into a hyper-realistic face portrait.", "lora": "Hyper-Realistic-Portrait"},
{"images": ["examples/A.jpeg"], "prompt": "Rotate the camera 45 degrees to the right.", "lora": "Multiple-Angles"},
{"images": ["examples/U.jpg"], "prompt": "Upscale this picture to 4K resolution.", "lora": "Upscaler"},
{"images": ["examples/L1.jpg", "examples/L2.jpg"], "prompt": "Apply the lighting from image 2 to image 1.", "lora": "Any-light"},
{"images": ["examples/PP1.jpg"], "prompt": "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed preserving realistic texture and details.", "lora": "Polaroid-Photo"},
{"images": ["examples/Z1.jpg"], "prompt": "Front-right quarter view.", "lora": "Fal-Multiple-Angles"},
{"images": ["examples/URP.jpg"], "prompt": "Transform into a cinematic flat log.", "lora": "Cinematic-FlatLog"},
{"images": ["examples/SL.jpg"], "prompt": "Neutral uniform lighting. Preserve identity and composition.", "lora": "Studio-DeLight"},
{"images": ["examples/PI.jpg"], "prompt": "Transform it into Pixar-inspired 3D.", "lora": "Pixar-Inspired-3D"},
{"images": ["examples/MT.jpg"], "prompt": "Paint with manga tone.", "lora": "Manga-Tone"},
{"images": ["examples/NCB.jpg"], "prompt": "Transform into a noir comic book style.", "lora": "Noir-Comic-Book"},
{"images": ["examples/URP.jpg"], "prompt": "Ultra-realistic portrait.", "lora": "Ultra-Realistic-Portrait"},
{"images": ["examples/MN.jpg"], "prompt": "Transform into Midnight Noir Eyes Spotlight.", "lora": "Midnight-Noir-Eyes-Spotlight"},
{"images": ["examples/ST1.jpg", "examples/ST2.jpg"], "prompt": "Convert Image 1 to the style of Image 2.", "lora": "Style-Transfer"},
{"images": ["examples/R1.jpg"], "prompt": "Change the picture to realistic photograph.", "lora": "Anything2Real"},
{"images": ["examples/UA.jpeg"], "prompt": "Unblur and upscale.", "lora": "Unblur-Anything"},
{"images": ["examples/L1.jpg", "examples/L2.jpg"], "prompt": "Refer to the color tone, remove the original lighting from Image 1, and relight Image 1 based on the lighting and color tone of Image 2.", "lora": "Light-Migration"},
{"images": ["examples/P1.jpg"], "prompt": "Transform into anime (while preserving the background and remaining elements maintaining realism and original details.)", "lora": "Anime-V2"},
]
def make_thumb_b64(path, max_dim=220):
if not os.path.exists(path):
return ""
try:
img = Image.open(path).convert("RGB")
img.thumbnail((max_dim, max_dim), LANCZOS)
buf = BytesIO()
img.save(buf, format="JPEG", quality=65)
return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
except Exception as e:
print(f"Thumbnail error for {path}: {e}")
return ""
def encode_full_image(path):
if not os.path.exists(path):
return ""
try:
with open(path, "rb") as f:
data = f.read()
ext = path.rsplit(".", 1)[-1].lower()
mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
return f"data:{mime};base64,{base64.b64encode(data).decode()}"
except Exception as e:
print(f"Encode error for {path}: {e}")
return ""
def build_example_cards_html():
cards = ""
for i, ex in enumerate(EXAMPLES_CONFIG):
thumbs_html = ""
for path in ex["images"]:
thumb = make_thumb_b64(path)
if thumb:
thumbs_html += f'<img src="{thumb}" alt="">'
else:
thumbs_html += '<div class="example-thumb-placeholder">Preview</div>'
n = len(ex["images"])
img_badge = f'{n} image{"s" if n > 1 else ""}'
lora_badge = html_lib.escape(ex["lora"])
prompt_short = html_lib.escape(ex["prompt"][:85])
if len(ex["prompt"]) > 85:
prompt_short += "…"
cards += f'''<div class="example-card" data-idx="{i}">
<div class="example-thumbs">{thumbs_html}</div>
<div class="example-meta">
<span class="example-badge">{img_badge}</span>
<span class="example-lora-badge">{lora_badge}</span>
</div>
<div class="example-prompt-text">{prompt_short}</div>
</div>'''
return cards
def load_example_data(idx_str):
try:
idx = int(float(idx_str)) if idx_str and idx_str.strip() else -1
except (ValueError, TypeError):
idx = -1
if idx < 0 or idx >= len(EXAMPLES_CONFIG):
return json.dumps({"images": [], "prompt": "", "lora": "", "names": [], "status": "error"})
ex = EXAMPLES_CONFIG[idx]
b64_list, names = [], []
for path in ex["images"]:
b64 = encode_full_image(path)
if b64:
b64_list.append(b64)
names.append(os.path.basename(path))
return json.dumps({"images": b64_list, "prompt": ex["prompt"], "lora": ex["lora"], "names": names, "status": "ok"})
print("Example gallery strip disabled (UI uses LoRA demos + chips instead).")
EXAMPLE_CARDS_HTML = ""
def b64_to_pil_list(b64_json_str):
if not b64_json_str or b64_json_str.strip() in ("", "[]"):
return []
try:
b64_list = json.loads(b64_json_str)
except Exception:
return []
pil_images = []
for b64_str in b64_list:
if not b64_str or not isinstance(b64_str, str):
continue
try:
if b64_str.startswith("data:image"):
_, data = b64_str.split(",", 1)
else:
data = b64_str
image_data = base64.b64decode(data)
pil_images.append(Image.open(BytesIO(image_data)).convert("RGB"))
except Exception as e:
print(f"Error decoding image: {e}")
return pil_images
# Long-side base sizes (normal = 1024, the previous default).
# "original" uses source pixels (long side capped at MAX_SOURCE_SIDE).
IMAGE_SIZE_BASE = {
"x-small": 512, # thumbnail
"small": 768,
"normal": 1024,
"large": 1280,
"x-large": 1536,
}
MAX_SOURCE_SIDE = 4000
def _round8(n: int) -> int:
return max(8, (int(n) // 8) * 8)
def _cap_source_dims(w: int, h: int):
"""Scale source so the long side is at most MAX_SOURCE_SIDE."""
long_side = max(w, h)
if long_side <= MAX_SOURCE_SIDE or long_side <= 0:
return w, h
scale = MAX_SOURCE_SIDE / float(long_side)
return max(1, int(round(w * scale))), max(1, int(round(h * scale)))
def update_dimensions_on_upload(image, aspect="original", size="normal"):
"""Compute generation width/height from source image, aspect preset, and size mode."""
if image is None:
base = IMAGE_SIZE_BASE.get(size, 1024) if size != "original" else 1024
return base, base
w, h = image.size
if size == "original":
# Exact source pixels (//8), long side capped at 4000.
# Non-original ratios: same long side, forced shape (F decision).
cw, ch = _cap_source_dims(w, h)
long_side = max(cw, ch)
if aspect == "square":
target_width = target_height = long_side
elif aspect == "wide":
target_width = long_side
target_height = int(long_side * 9 / 16)
elif aspect == "portrait":
target_height = long_side
target_width = int(long_side * 9 / 16)
else:
target_width, target_height = cw, ch
return _round8(target_width), _round8(target_height)
base = IMAGE_SIZE_BASE.get(size, 1024)
# Scale fixed aspect presets from their normal (1024-class) dimensions.
scale = base / 1024.0
if aspect == "square":
target_width = target_height = base
elif aspect == "wide":
# 16:9 at normal β‰ˆ 1344Γ—768
target_width = int(1344 * scale)
target_height = int(768 * scale)
elif aspect == "portrait":
# 9:16 at normal β‰ˆ 768Γ—1344
target_width = int(768 * scale)
target_height = int(1344 * scale)
else: # "original" β€” maintain aspect ratio, long side = base
if w > h:
target_width = base
target_height = int(target_width * h / w)
else:
target_height = base
target_width = int(target_height * w / h)
return _round8(target_width), _round8(target_height)
@spaces.GPU(size="xlarge")
def infer(
images_b64_json,
prompt,
lora_adapter,
seed,
randomize_seed,
guidance_scale,
steps,
aspect_ratio,
image_size,
prompt_pack="none",
progress=gr.Progress(track_tqdm=True),
):
gc.collect()
torch.cuda.empty_cache()
pil_images = b64_to_pil_list(images_b64_json)
if not pil_images:
raise gr.Error("Please upload at least one image to edit.")
if not prompt or prompt.strip() == "":
raise gr.Error("Please enter an edit prompt.")
# Prompt Packs: menu selection overrides prompt tags; tags work when menu is None
prompt, pil_images, lora_adapter, _applied_packs = apply_prompt_packs(
prompt, pil_images, lora_adapter, menu_pack_id=prompt_pack
)
if _applied_packs:
try:
gr.Info("Prompt packs: " + ", ".join(_applied_packs))
except Exception:
pass
if not lora_adapter or lora_adapter == "none":
print("--- No LoRA selected; running base model ---")
try:
pipe.disable_lora()
except Exception:
try:
pipe.set_adapters([])
except Exception as e:
print(f"Warning: could not clear adapters: {e}")
else:
spec = ADAPTER_SPECS.get(lora_adapter)
if not spec:
raise gr.Error(f"Configuration not found for: {lora_adapter}")
adapter_name = spec["adapter_name"]
if adapter_name not in LOADED_ADAPTERS:
print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
try:
pipe.load_lora_weights(spec["repo"], weight_name=spec["weights"], adapter_name=adapter_name)
LOADED_ADAPTERS.add(adapter_name)
except Exception as e:
raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
else:
print(f"--- Adapter {lora_adapter} already loaded. ---")
pipe.set_adapters([adapter_name], adapter_weights=[1.0])
if randomize_seed:
seed = random.randint(0, MAX_SEED)
generator = torch.Generator(device=device).manual_seed(seed)
width, height = update_dimensions_on_upload(pil_images[0], aspect_ratio, image_size)
# true_cfg_scale <= 1 disables classifier-free guidance; negative_prompt is then ignored
# and diffusers logs a warning. Only pass it when CFG is actually enabled (guidance > 1).
pipe_kwargs = dict(
image=pil_images,
prompt=prompt,
height=height,
width=width,
num_inference_steps=steps,
generator=generator,
true_cfg_scale=guidance_scale,
)
if guidance_scale is not None and float(guidance_scale) > 1.0:
pipe_kwargs["negative_prompt"] = (
"worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, "
"extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
)
try:
result_image = pipe(**pipe_kwargs).images[0]
return result_image, seed
except Exception as e:
raise e
finally:
gc.collect()
torch.cuda.empty_cache()
# ── UI assets (CSS / JS / HTML live under ui/ so app.py stays small) ──────────
_UI_DIR = Path(__file__).resolve().parent / "ui"
def _load_ui(name: str) -> str:
return (_UI_DIR / name).read_text(encoding="utf-8")
def _fill_ui(template: str, mapping: dict) -> str:
out = template
for key, val in mapping.items():
out = out.replace(f"__{key}__", val)
return out
css = _load_ui("app.css") + "\n" + _load_ui("app_phone.css")
gallery_js = _load_ui("gallery.js")
wire_outputs_js = _load_ui("wire_outputs.js")
boot_js = _load_ui("boot.js")
# ── SVG assets (same family/style as the FireRed-Image-Edit app, recolored via CSS) ──
FIRE_LOGO_SVG = '<svg viewBox="0 0 24 24" fill="white" xmlns="http://www.w3.org/2000/svg"><path d="M12 23c-3.6 0-8-2.69-8-7.5 0-3.5 3-6.5 4.5-8 .27-.27.75-.08.75.28v2.44c0 .42.5.63.72.28C12.28 7.5 13 3 13 1c0-.42.48-.64.8-.35C18 4.5 20 9 20 12c0 5.5-3.5 11-8 11z"/></svg>'
DOWNLOAD_SVG = '<svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M12 16l-5-5h3V4h4v7h3l-5 5z"/><path d="M20 18H4v2h16v-2z"/></svg>'
UPLOAD_SVG = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4"/><polyline points="17 8 12 3 7 8"/><line x1="12" y1="3" x2="12" y2="15"/></svg>'
REMOVE_SVG = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="12" cy="12" r="10"/><line x1="15" y1="9" x2="9" y2="15"/><line x1="9" y1="9" x2="15" y2="15"/></svg>'
CLEAR_SVG = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="3 6 5 6 21 6"/><path d="M19 6v14a2 2 0 01-2 2H7a2 2 0 01-2-2V6m3 0V4a2 2 0 012-2h4a2 2 0 012 2v2"/><line x1="10" y1="11" x2="10" y2="17"/><line x1="14" y1="11" x2="14" y2="17"/></svg>'
GITHUB_SVG = '<svg width="15" height="15" viewBox="0 0 16 16" xmlns="http://www.w3.org/2000/svg"><path fill="#ffffff" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0016 8c0-4.42-3.58-8-8-8z"/></svg>'
LORA_OPTIONS_HTML = "\n".join(
['<option value="none" selected>Choose a LoRA…</option>']
+ [
f'<option value="{html_lib.escape(name)}">{html_lib.escape(name)}</option>'
for name in LORA_DISPLAY_ORDER
]
)
# ── Gradio app ─────────────────────────────────────────────────────────────────
with gr.Blocks() as demo:
hidden_images_b64 = gr.Textbox(value="[]", elem_id="hidden-images-b64", elem_classes="hidden-input", container=False)
prompt = gr.Textbox(value="", elem_id="prompt-gradio-input", elem_classes="hidden-input", container=False)
lora_adapter = gr.Dropdown(choices=["none"] + list(LORA_DISPLAY_ORDER), value="none", elem_id="gradio-lora", elem_classes="hidden-input", container=False)
prompt_pack = gr.Dropdown(choices=["none"] + list(PROMPT_PACK_IDS), value="none", elem_id="gradio-prompt-pack", elem_classes="hidden-input", container=False)
seed = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=0, elem_id="gradio-seed", elem_classes="hidden-input", container=False)
randomize_seed = gr.Checkbox(value=True, elem_id="gradio-randomize", elem_classes="hidden-input", container=False)
guidance_scale = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=1.0, elem_id="gradio-guidance", elem_classes="hidden-input", container=False)
steps = gr.Slider(minimum=1, maximum=50, step=1, value=4, elem_id="gradio-steps", elem_classes="hidden-input", container=False)
result = gr.Image(elem_id="gradio-result", elem_classes="hidden-input", container=False, format="png")
example_idx = gr.Textbox(value="", elem_id="example-idx-input", elem_classes="hidden-input", container=False)
example_result = gr.Textbox(value="", elem_id="example-result-data", elem_classes="hidden-input", container=False)
example_load_btn = gr.Button("Load Example", elem_id="example-load-btn")
aspect_ratio = gr.Dropdown(
choices=[
("Original (maintain aspect ratio)", "original"),
("Square (1:1)", "square"),
("Wide (16:9)", "wide"),
("Portrait (9:16)", "portrait"),
],
value="original",
elem_id="gradio-aspect",
elem_classes="hidden-input",
container=False,
)
image_size = gr.Dropdown(
choices=[
("Original (source pixels)", "original"),
("X-Small (thumbnail)", "x-small"),
("Small", "small"),
("Normal", "normal"),
("Large", "large"),
("X-Large", "x-large"),
],
value="normal",
elem_id="gradio-size",
elem_classes="hidden-input",
container=False,
)
_shell_fill = {
"FIRE_LOGO_SVG": FIRE_LOGO_SVG,
"GITHUB_SVG": GITHUB_SVG,
"UPLOAD_SVG": UPLOAD_SVG,
"REMOVE_SVG": REMOVE_SVG,
"CLEAR_SVG": CLEAR_SVG,
"DOWNLOAD_SVG": DOWNLOAD_SVG,
"LORA_OPTIONS_HTML": LORA_OPTIONS_HTML,
"PROMPT_PACK_OPTIONS_HTML": PROMPT_PACK_OPTIONS_HTML,
}
shell_desktop_html = _fill_ui(_load_ui("shell.html"), _shell_fill)
shell_phone_html = _fill_ui(_load_ui("shell_phone.html"), _shell_fill)
# Both skins in DOM briefly; boot.js removes the unused one so element IDs stay unique.
# Phone stack: upload β†’ output β†’ prompts β†’ prompt β†’ menus β†’ LoRA β†’ advanced β†’ Edit β†’ footer
dual_shell_html = (
'<div id="qwen-ui-root">'
f'<div id="skin-desktop" data-skin="desktop">{shell_desktop_html}</div>'
f'<div id="skin-phone" data-skin="phone" hidden>{shell_phone_html}</div>'
"</div>"
)
gr.HTML(dual_shell_html)
# LoRA help + demo thumbs: Gradio 6 won't run scripts inside gr.HTML, so use demo.load.
lora_demo_js = _fill_ui(
_load_ui("lora_demo.js"),
{
"LORA_HELP_JS": LORA_HELP_JS,
"LORA_DEFAULT_PROMPTS_JS": LORA_DEFAULT_PROMPTS_JS,
"LORA_DEMO_BEFORE_JS": LORA_DEMO_BEFORE_JS,
"LORA_DEMO_AFTER_JS": LORA_DEMO_AFTER_JS,
"LORA_DEMO_GENERIC_BEFORE_JS": LORA_DEMO_GENERIC_BEFORE_JS,
"LORA_DEMO_GENERIC_AFTER_JS": LORA_DEMO_GENERIC_AFTER_JS,
},
)
run_btn = gr.Button("Run", elem_id="gradio-run-btn")
# boot first: pick skin before gallery/lora_demo bind to IDs
demo.load(fn=None, js=boot_js)
demo.load(fn=None, js=lora_demo_js)
demo.load(fn=None, js=gallery_js)
demo.load(fn=None, js=wire_outputs_js)
run_btn.click(
fn=infer,
inputs=[hidden_images_b64, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps, aspect_ratio, image_size, prompt_pack],
outputs=[result, seed],
js=r"""(imgs, p, la, s, rs, gs, st, ar, sz, pp) => {
const images = window.__uploadedImages || [];
const b64Array = images.map(img => img.b64);
const imgsJson = JSON.stringify(b64Array);
const promptEl = document.getElementById('custom-prompt-input');
const loraEl = document.getElementById('custom-lora-select');
const packEl = document.getElementById('custom-prompt-pack-select');
const aspectEl = document.getElementById('custom-aspect-select');
const sizeEl = document.getElementById('custom-size-select');
const promptVal = promptEl ? promptEl.value : p;
const loraVal = loraEl ? loraEl.value : la;
const packVal = packEl ? packEl.value : (pp || 'none');
const aspectVal = aspectEl ? aspectEl.value : (ar || 'original');
const sizeVal = sizeEl ? sizeEl.value : (sz || 'normal');
return [imgsJson, promptVal, loraVal, s, rs, gs, st, aspectVal, sizeVal, packVal];
}""",
)
example_load_btn.click(
fn=load_example_data,
inputs=[example_idx],
outputs=[example_result],
queue=False,
)
if __name__ == "__main__":
_app_root = Path(__file__).resolve().parent
demo.queue(max_size=50).launch(
css=css,
mcp_server=True,
ssr_mode=False,
show_error=True,
# Absolute paths so /gradio_api/file=<abs> demo thumbs work on Spaces
allowed_paths=[
str(_app_root / "examples"),
str(_app_root / "loras"),
str(_app_root / "prompt_packs"),
],
)