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)),
"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"(? 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 and optional ref as image #2.
Menu selection (menu_pack_id) overrides prompt triggers when not "none".
Does not change the LoRA — that stays whatever the user selected (default 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}")
return new_prompt, new_images, lora_adapter, 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_option_html(p: dict) -> str:
vid = html_lib.escape(p["id"], quote=True)
label = p["name"] + (f" · {p['id']}" if p["name"] != p["id"] else "")
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)."
)
desc_attr = html_lib.escape(desc, quote=True)
return f''
PROMPT_PACK_OPTIONS_HTML = "\n".join(
[
''.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''
else:
thumbs_html += '
Preview
'
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'''
{thumbs_html}
{img_badge}{lora_badge}
{prompt_short}
'''
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 = ''
DOWNLOAD_SVG = ''
UPLOAD_SVG = ''
REMOVE_SVG = ''
CLEAR_SVG = ''
GITHUB_SVG = ''
LORA_OPTIONS_HTML = "\n".join(
['']
+ [
f''
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 = (
'