inference.py

from __future__ import annotations

import base64
import hashlib
import io
import json
import math
import os
import random
import re
import sys
import contextlib
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import requests
import torch
import torchvision.transforms.functional as TF
from huggingface_hub import snapshot_download
from PIL import Image

from preprocess import select_primary_garment


DEFAULT_NEG_PROMPT = (
    "low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, "
    "poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, "
    "bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, "
    "missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers."
)

DEFAULT_PROMPTS = {
    "upper_body": "Make the person in Picture 2 wear the top from Picture 1. Keep the pose, body shape, background, and lighting from Picture 2.",
    "lower_body": "Make the person in Picture 2 wear the bottom from Picture 1. Keep the pose, body shape, background, and lighting from Picture 2.",
    "dress": "Make the person in Picture 2 wear the dress from Picture 1. Keep the pose, body shape, background, and lighting from Picture 2.",
    "full_body": "Make the person in Picture 2 wear the clothing from Picture 1. Keep the pose, body shape, background, and lighting from Picture 2.",
}

OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://modelservice.jdcloud.com/v1")
QWEN235_MODEL = os.getenv("QWEN235_MODEL", "qwen3-vl-235")
QWEN235_CHAT_URL = os.getenv("QWEN235_CHAT_URL") or f"{OPENAI_BASE_URL.rstrip('/')}/chat/completions"
PE_SYSTEM_ROLE = "You are an expert vision AI prompt engineer."
PE_TEMPLATE_DEFAULT = (
    "# Multi-Image Edit Instruction Rewriter\n\n"
    "You are an expert at refining MULTI-IMAGE editing instructions. Given the "
    "user's original prompt and the uploaded reference images (in order: Image 1, "
    "Image 2, Image 3, Image 4), produce a single improved English editing "
    "caption that a diffusion-based multi-image editing model can follow precisely.\n\n"

    "## Core Rules\n"
    "1. **Preserve intent.** Never invent new operations. Only clarify, enrich "
    "and correct what the user already asked for.\n"
    "2. **Canonical image references.** Always refer to the inputs as "
    "'Image 1 / Image 2 / Image 3 / Image 4'. Rewrite vague mentions like "
    "'图1 / 第一张 / img1 / the first one' into the canonical form.\n"
    "3. **Source-of-truth attribution.** For every visual element in the output, "
    "state explicitly WHICH reference image it comes from "
    "(e.g. 'the person from Image 1', 'the dress from Image 2').\n"
    "4. **Default role mapping.** If the role of each reference image is "
    "ambiguous, assume Image 1 is the SOURCE / main subject image and the "
    "remaining images provide assets (garments, objects, faces, scenes, styles).\n"
    "5. **Output the rewritten English caption only.** No preamble, no JSON, "
    "no markdown fences, no explanations.\n\n"

    "## Identify the Edit Type First (then apply the matching rules)\n"
    "Before writing the caption, internally classify the user instruction into "
    "ONE of the following categories. The category determines how much you may "
    "elaborate and what to lock down.\n\n"

    "### A. Replacement / Virtual Try-On  —  *consistency first*\n"
    "Trigger words: 'replace', 'swap', 'try on', '换', '穿', '把X换成Y', "
    "background swap, face swap, product swap.\n"
    "- **Golden rule: CONSISTENCY OVER CREATIVITY.** This is NOT artistic "
    "generation. Do NOT add cinematic effects, extra props, new lighting, new "
    "camera angles or any creative embellishments.\n"
    "- Preserve EVERYTHING in the source image (Image 1) that is not explicitly "
    "being replaced: identity, face, hair, skin tone, body shape, pose, "
    "expression, background, lighting, color tone, camera angle, composition.\n"
    "- **Try-On template:** `Keep the person in Image 1 strictly unchanged "
    "(identity, face, hair, pose, body shape, background, lighting). Replace "
    "ONLY the <garment type> with the <garment> from Image N, preserving its "
    "color, pattern, texture, logo and silhouette. Adapt naturally to the "
    "person's pose with realistic fabric drape.`\n"
    "- **Object swap template:** `Replace <old object> in Image 1 with "
    "<new object> from Image N. Keep position, scale, perspective, lighting "
    "and shadow consistent. Everything else in Image 1 remains unchanged.`\n"
    "- **Face / person swap template:** `Replace the face in Image 1 with the "
    "face from Image N. Preserve the original pose, body, clothing, lighting "
    "and background of Image 1. Transfer facial geometry, skin tone, hair from Image N.`\n"
    "- **Background swap template:** `Keep the foreground subject of Image 1 "
    "unchanged. Replace ONLY the background with the scene from Image N. Match "
    "color temperature so the subject blends naturally.`\n"
    "- Vocabulary: plain, factual. AVOID 'cinematic', 'dreamy', 'volumetric "
    "fog', 'golden hour'. Target length ~60-100 words.\n\n"

    "### B. Subject-Driven Composition  —  *IP consistency + natural integration*\n"
    "Trigger words: 'put X into Y', 'let X hold/wear/use Y', 'X in the scene of Y'.\n"
    "- Preserve identity-defining features of every named subject (face, hairstyle, "
    "body proportions, signature outfit/logo, product shape and branding).\n"
    "- The scene, lighting and interaction MAY be re-imagined so the composition "
    "looks natural and coherent.\n"
    "- Describe in this order: (1) subjects + their source image + identity traits "
    "to preserve, (2) the interaction/action, (3) the scene/background (note if "
    "it comes from a reference or is newly generated), (4) brief, tasteful "
    "lighting & camera notes so the subjects look embedded naturally — avoid "
    "over-stylization. Target length ~80-140 words.\n\n"

    "### C. Creative Multi-Image Composition  —  *aesthetics matter*\n"
    "Trigger words: 'combine', 'blend', 'compose', 'fuse', '融合', '合成', or any "
    "request for a new artistic scene that uses multiple references.\n"
    "- Anchor key elements: for each reference, state in one short clause what to "
    "inherit (subject identity / outfit / style / environment / props).\n"
    "- Beyond that, freely design lighting (volumetric fog, golden hour, rim "
    "light...), color palette, texture, camera angle, focal length and composition "
    "(rule of thirds, leading lines...) to make the final image visually stunning.\n"
    "- Add evocative mood & atmosphere adjectives; avoid generic words like "
    "'good' or 'nice'. Target length ~120-200 words.\n\n"

    "### D. Reference-Based Style Transfer  —  *content from source, style from reference*\n"
    "Trigger words: 'in the style of', 'redraw ... like', '风格', 'apply the look of'.\n"
    "- Lock the content of the source image (Image 1): subjects, layout, pose, "
    "object positions, composition.\n"
    "- Summarize the target style (Image N) with 3-5 distinctive visual cues — "
    "palette, lighting, texture, era, medium (e.g. 'oil painting with thick "
    "impasto strokes, muted earth tones, soft rembrandt lighting, "
    "19th-century romanticism').\n"
    "- Explicitly say 'apply the style of Image N to the content of Image 1'. "
    "Target length ~80-120 words.\n\n"

    "### E. Image Combination / Layout  —  *preserve each image, arrange by layout*\n"
    "Trigger words: 'collage', 'grid', 'side by side', 'before/after', '九宫格', "
    "'拼图', '排版'.\n"
    "- Each reference image stays largely intact; only minor cropping/resizing "
    "for layout fit is allowed. Do NOT blend or restyle across images.\n"
    "- State the layout explicitly ('2x2 grid', 'left/right split', "
    "'three vertical panels', 'before/after horizontally'...).\n"
    "- Slot mapping: for each slot/position, name which reference image goes "
    "there (e.g. 'top-left: Image 1, top-right: Image 2, bottom: Image 3').\n"
    "- If text overlays are required, wrap text content in English double quotes "
    "`\"...\"`, preserving the original language and casing. Default to thin "
    "white gutters and clean neutral background if not specified. "
    "Target length ~60-100 words.\n\n"

    "## Sanity Checks\n"
    "- If the instruction contains contradictions, resolve them with the most "
    "reasonable interpretation.\n"
    "- Do not invent details that are not visible in the references or required "
    "by the user instruction.\n"
    "- If spatial placement is missing in a composition task, pick a "
    "compositionally balanced location (near the subject, in negative space, "
    "along rule-of-thirds lines, etc.).\n\n"

    "## Output Format\n"
    "Output ONLY the final English caption. No preamble, no JSON, no explanation."
)


@dataclass
class LoadedTryOnModel:
    pipeline: Any
    cfg: Any
    device: torch.device
    model_dir: Path


_MODEL: LoadedTryOnModel | None = None
_MODEL_DIR: Path | None = None
_MODEL_LOCK = threading.Lock()
_PE_CACHE: dict[str, str] = {}


def _log_stage(message: str) -> None:
    print(f"[tryon] {time.strftime('%Y-%m-%d %H:%M:%S')} | {message}", flush=True)


@contextlib.contextmanager
def _stage_heartbeat(stage: str, interval_s: int = 30):
    stop = threading.Event()

    def beat() -> None:
        started = time.time()
        while not stop.wait(interval_s):
            _log_stage(f"{stage} still running after {time.time() - started:.0f}s")

    thread = threading.Thread(target=beat, daemon=True)
    thread.start()
    try:
        yield
    finally:
        stop.set()
        thread.join(timeout=1)


def _get_bool_env(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _device() -> torch.device:
    if not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA is not available. Make sure the Space hardware is set to ZeroGPU "
            "and that inference is called from a function decorated with @spaces.GPU."
        )
    return torch.device("cuda")


def set_seed(seed: int) -> torch.Generator:
    seed = int(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    return torch.Generator(device="cuda").manual_seed(seed)


def _encode_image_base64_png(pil_image: Image.Image) -> str:
    buf = io.BytesIO()
    pil_image.convert("RGB").save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def _remove_think_block(text: str) -> str:
    return re.sub(r"<think>.*?</think>", "", text or "", flags=re.IGNORECASE | re.DOTALL).strip()


def _normalize_chat_message_content(content: Any) -> str:
    if content is None:
        return ""
    if isinstance(content, str):
        return content.strip()

    texts: list[str] = []

    def walk(value: Any) -> None:
        if value is None:
            return
        if isinstance(value, str):
            if value.strip():
                texts.append(value.strip())
            return
        if isinstance(value, dict):
            text = value.get("text")
            if isinstance(text, str) and text.strip():
                texts.append(text.strip())
            for child in value.values():
                if isinstance(child, (dict, list)):
                    walk(child)
            return
        if isinstance(value, list):
            for item in value:
                walk(item)

    walk(content)
    return "\n".join(texts).strip()


def _extract_caption_text(text: str) -> str:
    clean = _remove_think_block((text or "").strip())
    if not clean:
        return ""
    clean = clean.replace("```json", "").replace("```", "").strip()

    try:
        obj = json.loads(clean)
        if isinstance(obj, dict) and isinstance(obj.get("Rewritten"), str):
            return obj["Rewritten"].strip().replace("\n", " ")
    except Exception:
        pass

    first_l = clean.find("{")
    last_r = clean.rfind("}")
    if 0 <= first_l < last_r:
        candidate = clean[first_l:last_r + 1]
        try:
            obj = json.loads(candidate)
            if isinstance(obj, dict) and isinstance(obj.get("Rewritten"), str):
                return obj["Rewritten"].strip().replace("\n", " ")
        except Exception:
            pass

    return clean


def _image_signature(image: Image.Image | None) -> str:
    if image is None:
        return "none"
    try:
        buf = io.BytesIO()
        image.convert("RGB").save(buf, format="PNG")
        data = buf.getvalue()
        digest = hashlib.md5(data[:65536]).hexdigest()
        return f"{image.size[0]}x{image.size[1]}_{digest}"
    except Exception:
        return f"{image.size[0]}x{image.size[1]}"


def _images_signature(images: list[Image.Image]) -> str:
    if not images:
        return "no_images"
    return "|".join(_image_signature(img) for img in images)


def _normalize_images_input(images: list[Image.Image] | None) -> list[Image.Image]:
    out: list[Image.Image] = []
    for image in images or []:
        if image is None:
            continue
        try:
            out.append(image.convert("RGB"))
        except Exception:
            out.append(image)
    return out


def _get_pe_instruction_template() -> str:
    return (os.getenv("PE_INSTRUCTION_TEMPLATE") or PE_TEMPLATE_DEFAULT).strip()


def _build_pe_prompt_text(original_prompt: str, instruction_template: str) -> str:
    original_prompt = str(original_prompt or "").strip()
    instruction_template = str(instruction_template or "").strip()
    return (
        f"Based on the uploaded image(s) and the prompt: '{original_prompt}', "
        "please follow the instructions below to produce the rewritten caption.\n\n"
        f"{instruction_template}"
    )


def _build_openai_style_messages(prompt: str, images: list[Image.Image] | None, instruction_template: str) -> list[dict[str, Any]]:
    user_content: list[dict[str, Any]] = []
    for image in _normalize_images_input(images):
        user_content.append(
            {
                "type": "image_url",
                "image_url": {"url": f"data:image/png;base64,{_encode_image_base64_png(image)}"},
            }
        )
    user_content.append({"type": "text", "text": _build_pe_prompt_text(prompt, instruction_template)})
    return [
        {"role": "system", "content": PE_SYSTEM_ROLE},
        {"role": "user", "content": user_content},
    ]


def _rewrite_prompt_qwen3vl235_chatapi(
    prompt: str,
    images: list[Image.Image] | None,
    *,
    instruction_template: str | None = None,
    model: str | None = None,
    api_key: str | None = None,
    chat_url: str | None = None,
    max_retries: int | None = None,
    timeout_s: float | None = None,
    max_tokens: int | None = None,
) -> str:
    prompt = str(prompt or "").strip()
    if not prompt:
        return prompt

    api_key = api_key or os.getenv("QWEN235_API_KEY") or os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise EnvironmentError("QWEN235_API_KEY or OPENAI_API_KEY is not set")

    chat_url = chat_url or QWEN235_CHAT_URL
    if not chat_url:
        raise EnvironmentError("QWEN235_CHAT_URL is empty")

    payload = {
        "model": model or QWEN235_MODEL,
        "messages": _build_openai_style_messages(prompt, images, instruction_template or _get_pe_instruction_template()),
        "max_tokens": int(max_tokens or os.getenv("QWEN235_MAX_TOKENS", "1024")),
        "stream": False,
    }
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "Trace-Id": f"tryon-space-qwen235-{int(time.time() * 1000)}",
    }

    last_err: Exception | None = None
    retries = int(max_retries or os.getenv("QWEN235_MAX_RETRIES", "3"))
    timeout = float(timeout_s or os.getenv("QWEN235_TIMEOUT", "60"))
    for attempt in range(retries):
        try:
            resp = requests.post(chat_url, headers=headers, json=payload, timeout=timeout)
            if resp.status_code != 200:
                raise RuntimeError(f"Qwen3-VL-235 API failed: status={resp.status_code}, text={resp.text[:2000]}")
            data = resp.json()
            content = (((data.get("choices") or [{}])[0].get("message") or {}).get("content"))
            output_text = _normalize_chat_message_content(content)
            print(f"[Qwen235 raw] {repr(output_text[:500])}", flush=True)
            caption = _extract_caption_text(output_text)
            return caption or prompt
        except Exception as exc:
            last_err = exc
            time.sleep(0.5 * (2**attempt))

    raise RuntimeError(f"Prompt Enhancer failed: {last_err}")


def rewrite_prompt_if_enabled(prompt: str, images: list[Image.Image], enabled: bool) -> tuple[str, str]:
    prompt = str(prompt or "")
    valid_images = _normalize_images_input(images)
    if not enabled:
        return prompt, "Prompt Enhancer disabled."
    if not valid_images:
        return prompt, "Prompt Enhancer skipped: no reference images."

    instruction_template = _get_pe_instruction_template()
    template_signature = hashlib.md5(instruction_template.encode("utf-8")).hexdigest()[:12]
    cache_key = f"qwen235|edit_caption|p={prompt.strip()}|imgs={_images_signature(valid_images)}|tpl={template_signature}"
    if cache_key in _PE_CACHE:
        return _PE_CACHE[cache_key], "Prompt Enhancer cache hit."

    try:
        rewritten = _rewrite_prompt_qwen3vl235_chatapi(prompt, valid_images, instruction_template=instruction_template)
        _PE_CACHE[cache_key] = rewritten
        return rewritten, "Prompt Enhancer applied."
    except Exception as exc:
        return prompt, f"Prompt Enhancer failed, using original prompt: {exc}"


def _download_model_repo() -> Path:
    local_dir = os.getenv("MODEL_LOCAL_DIR")
    if local_dir:
        local_path = Path(local_dir)
        if not local_path.exists():
            raise FileNotFoundError(f"MODEL_LOCAL_DIR does not exist: {local_path}")
        return local_path

    model_id = os.getenv("MODEL_ID")
    if not model_id:
        raise RuntimeError(
            "MODEL_ID is not set. Add it as a Hugging Face Space Variable, or set "
            "MODEL_LOCAL_DIR to a mounted bucket/model directory."
        )

    token = os.getenv("HF_TOKEN")
    revision = os.getenv("MODEL_REVISION") or None

    kwargs = {
        "repo_id": model_id,
        "revision": revision,
        "token": token,
    }

    return Path(snapshot_download(**kwargs))


def _add_code_paths(model_dir: Path) -> None:
    candidates = [
        Path(os.getenv("JOY_CODE_DIR", "")) if os.getenv("JOY_CODE_DIR") else None,
        model_dir,
        model_dir / "code",
        model_dir / "JoyAI-Image",
        model_dir / "JoyAI-Image-a4afa71",
    ]
    candidates.extend(sorted(path for path in model_dir.glob("JoyAI-Image*") if path.is_dir()))
    for candidate in candidates:
        if candidate and candidate.exists() and str(candidate) not in sys.path:
            sys.path.insert(0, str(candidate))


def _resolve_existing_path(value: str | None, model_dir: Path, *, required: bool, description: str) -> Path | None:
    if not value:
        if required:
            raise RuntimeError(f"{description} is required but was not provided.")
        return None

    path = Path(value)
    if not path.is_absolute():
        path = model_dir / path
    if required and not path.exists():
        raise FileNotFoundError(f"{description} not found: {path}")
    return path


def _find_default_config(model_dir: Path) -> Path:
    env_path = _resolve_existing_path(os.getenv("JOY_CONFIG_PATH"), model_dir, required=False, description="JOY_CONFIG_PATH")
    if env_path:
        return env_path

    candidates = [
        model_dir / "configs" / "editing_exp" / "edit_1024p_demo.py",
        model_dir / "JoyAI-Image-a4afa71" / "configs" / "editing_exp" / "edit_1024p_demo.py",
        model_dir / "edit_1024p_demo.py",
        model_dir / "config.py",
        model_dir / "configs" / "hf_space_config.py",
    ]
    candidates.extend(
        sorted(path / "configs" / "editing_exp" / "edit_1024p_demo.py" for path in model_dir.glob("JoyAI-Image*") if path.is_dir())
    )
    for candidate in candidates:
        if candidate.exists():
            return candidate
    raise FileNotFoundError(
        "Could not find a JoyAI config file. Set JOY_CONFIG_PATH to a config .py path "
        "inside the model repo or Space."
    )


def _find_transformer_path(model_dir: Path) -> Path:
    env_path = _resolve_existing_path(
        os.getenv("JOY_TRANSFORMER_PATH") or os.getenv("JOY_DIT_CKPT"),
        model_dir,
        required=False,
        description="JOY_TRANSFORMER_PATH",
    )
    if env_path:
        return env_path

    transformer_dir = model_dir / "transformer"
    if transformer_dir.exists():
        return transformer_dir

    preferred_names = ["step_33000.pth", "step_25500.pth", "transformer.pth", "dit.pth", "model.pth"]
    for name in preferred_names:
        matches = list(model_dir.rglob(name))
        if matches:
            return matches[0]

    matches = list(model_dir.rglob("*.pth"))
    if matches:
        return matches[0]

    raise FileNotFoundError(
        "Could not find a transformer .pth checkpoint in the model repo. "
        "Set JOY_TRANSFORMER_PATH or JOY_DIT_CKPT."
    )


def _patch_config_paths(cfg: Any, model_dir: Path) -> None:
    vae_path = os.getenv("JOY_VAE_PATH")
    if hasattr(cfg, "vae_arch_config"):
        resolved = _resolve_existing_path(vae_path, model_dir, required=True, description="JOY_VAE_PATH") if vae_path else _resolve_default_vae_path(model_dir)
        if resolved is None:
            return
        cfg.vae_arch_config["pretrained"] = str(resolved)

    text_encoder_path = _resolve_text_encoder_path(model_dir)
    if text_encoder_path and hasattr(cfg, "text_encoder_arch_config"):
        params = cfg.text_encoder_arch_config.setdefault("params", {})
        params["text_encoder_ckpt"] = str(text_encoder_path)


def _has_text_encoder_weights(path: Path) -> bool:
    patterns = ["pytorch_model.bin", "model.safetensors", "model-*.safetensors", "pytorch_model-*.bin"]
    return any(list(path.glob(pattern)) for pattern in patterns)


def _resolve_model_root(model_dir: Path) -> Path | None:
    value = os.getenv("JOY_MODEL_PATH") or os.getenv("JOY_BASE_MODEL_PATH")
    if value:
        return _resolve_existing_path(value, model_dir, required=True, description="JOY_MODEL_PATH")

    candidates = [
        model_dir,
        model_dir / "FireRed-Image-Edit-1.0",
        model_dir / "model_hub" / "FireRed-Image-Edit-1.0",
    ]
    for candidate in candidates:
        if candidate.exists() and any((candidate / name).exists() for name in ("text_encoder", "processor", "tokenizer", "vae")):
            return candidate
    return None


def _resolve_default_vae_path(model_dir: Path) -> Path | None:
    model_root = _resolve_model_root(model_dir)
    candidates = [
        model_dir / "vae" / "diffusion_pytorch_model.safetensors",
        model_dir / "vae" / "Wan2.1_VAE.pth",
    ]
    if model_root is not None:
        candidates.extend(
            [
                model_root / "vae" / "diffusion_pytorch_model.safetensors",
                model_root / "vae" / "Wan2.1_VAE.pth",
            ]
        )
    for candidate in candidates:
        if candidate.exists():
            return candidate
    return None


def _resolve_text_encoder_path(model_dir: Path) -> Path | None:
    value = os.getenv("JOY_TEXT_ENCODER_PATH")
    if value:
        return _resolve_existing_path(value, model_dir, required=True, description="JOY_TEXT_ENCODER_PATH")

    model_root = _resolve_model_root(model_dir)
    if model_root is not None:
        text_encoder_dir = model_root / "text_encoder"
        if text_encoder_dir.is_dir() and _has_text_encoder_weights(text_encoder_dir):
            return text_encoder_dir
        if _has_text_encoder_weights(model_root):
            return model_root
    return None


def _resolve_processor_path(model_dir: Path) -> Path | None:
    value = os.getenv("JOY_PROCESSOR_PATH")
    if value:
        return _resolve_existing_path(value, model_dir, required=True, description="JOY_PROCESSOR_PATH")

    model_root = _resolve_model_root(model_dir)
    if model_root is not None and (model_root / "processor").is_dir():
        return model_root / "processor"
    return None


def _resolve_tokenizer_path(model_dir: Path) -> Path | None:
    value = os.getenv("JOY_TOKENIZER_PATH")
    if value:
        return _resolve_existing_path(value, model_dir, required=True, description="JOY_TOKENIZER_PATH")

    model_root = _resolve_model_root(model_dir)
    if model_root is not None and (model_root / "tokenizer").is_dir():
        return model_root / "tokenizer"
    return None


@contextlib.contextmanager
def _patch_split_firered_paths(
    *,
    text_encoder_path: Path | None,
    processor_path: Path | None,
    tokenizer_path: Path | None,
):
    if text_encoder_path is None or (processor_path is None and tokenizer_path is None):
        yield
        return

    import transformers
    import xvideo.models.common.diffusion.pipelines.pipeline as pipeline_module
    import xvideo.models.mmdit.text_encoder as text_encoder_module

    original_auto_processor = transformers.AutoProcessor.from_pretrained
    original_pipeline_auto_processor = pipeline_module.AutoProcessor.from_pretrained
    original_qwen_tokenizer = transformers.Qwen2Tokenizer.from_pretrained
    original_text_encoder_qwen_tokenizer = text_encoder_module.Qwen2Tokenizer.from_pretrained

    text_encoder_str = str(text_encoder_path)

    def patched_auto_processor_from_pretrained(pretrained_model_name_or_path, *args, **kwargs):
        if processor_path is not None and str(pretrained_model_name_or_path) == text_encoder_str:
            return original_auto_processor(str(processor_path), *args, **kwargs)
        return original_auto_processor(pretrained_model_name_or_path, *args, **kwargs)

    def patched_qwen_tokenizer_from_pretrained(pretrained_model_name_or_path, *args, **kwargs):
        if tokenizer_path is not None and str(pretrained_model_name_or_path) == text_encoder_str:
            return original_qwen_tokenizer(str(tokenizer_path), *args, **kwargs)
        return original_qwen_tokenizer(pretrained_model_name_or_path, *args, **kwargs)

    transformers.AutoProcessor.from_pretrained = patched_auto_processor_from_pretrained
    pipeline_module.AutoProcessor.from_pretrained = patched_auto_processor_from_pretrained
    transformers.Qwen2Tokenizer.from_pretrained = patched_qwen_tokenizer_from_pretrained
    text_encoder_module.Qwen2Tokenizer.from_pretrained = patched_qwen_tokenizer_from_pretrained
    try:
        yield
    finally:
        transformers.AutoProcessor.from_pretrained = original_auto_processor
        pipeline_module.AutoProcessor.from_pretrained = original_pipeline_auto_processor
        transformers.Qwen2Tokenizer.from_pretrained = original_qwen_tokenizer
        text_encoder_module.Qwen2Tokenizer.from_pretrained = original_text_encoder_qwen_tokenizer


def load_model() -> LoadedTryOnModel:
    """Load the JoyAI try-on model once inside the ZeroGPU-allocated process."""
    global _MODEL, _MODEL_DIR

    if _MODEL is not None:
        return _MODEL

    with _MODEL_LOCK:
        if _MODEL is not None:
            return _MODEL

        started_at = time.time()
        _log_stage("load_model start")
        _MODEL_DIR = _download_model_repo()
        _add_code_paths(_MODEL_DIR)

        from xvideo.config import load_config_class_from_pyfile
        from xvideo.models import load_dit, load_pipeline

        config_path = _find_default_config(_MODEL_DIR)
        transformer_path = _find_transformer_path(_MODEL_DIR)
        text_encoder_path = _resolve_text_encoder_path(_MODEL_DIR)
        processor_path = _resolve_processor_path(_MODEL_DIR)
        tokenizer_path = _resolve_tokenizer_path(_MODEL_DIR)
        _log_stage(f"resolved config={config_path}")
        _log_stage(f"resolved transformer={transformer_path}")
        _log_stage(f"resolved text_encoder={text_encoder_path}")
        _log_stage(f"resolved processor={processor_path}")
        _log_stage(f"resolved tokenizer={tokenizer_path}")

        config_class = load_config_class_from_pyfile(str(config_path))
        cfg = config_class()
        cfg.dit_ckpt = str(transformer_path)
        cfg.dit_ckpt_type = os.getenv(
            "JOY_DIT_CKPT_TYPE",
            "safetensor" if transformer_path.is_dir() else getattr(cfg, "dit_ckpt_type", "pt"),
        )
        cfg.training_mode = False
        cfg.use_fsdp_inference = False
        cfg.hsdp_shard_dim = 1
        _patch_config_paths(cfg, _MODEL_DIR)

        device = _device()
        _log_stage(f"loading DiT on {device}")
        with _patch_split_firered_paths(
            text_encoder_path=text_encoder_path,
            processor_path=processor_path,
            tokenizer_path=tokenizer_path,
        ):
            dit_started_at = time.time()
            with _stage_heartbeat("loading DiT"):
                dit = load_dit(cfg, device=device)
            _log_stage(f"loaded DiT in {time.time() - dit_started_at:.1f}s")
            _log_stage("freezing DiT parameters")
            dit.requires_grad_(False)
            _log_stage("setting DiT eval mode")
            dit.eval()
            pipeline_started_at = time.time()
            _log_stage("loading pipeline")
            with _stage_heartbeat("loading pipeline"):
                pipeline = load_pipeline(cfg, dit, device)
            _log_stage(f"loaded pipeline in {time.time() - pipeline_started_at:.1f}s")

        _MODEL = LoadedTryOnModel(pipeline=pipeline, cfg=cfg, device=device, model_dir=_MODEL_DIR)
        _log_stage(f"load_model done in {time.time() - started_at:.1f}s")
        return _MODEL


def get_bucket_size(img: Image.Image) -> tuple[int, int]:
    from xvideo.config import generate_video_image_bucket
    from xvideo.dataset.bucket_util import BucketGroup

    bucket_configs = generate_video_image_bucket(
        basesize=1024,
        min_temporal=1,
        max_temporal=1,
        bs_img=8,
        bs_vid=4,
        bs_mimg=8,
        min_items=1,
        max_items=6,
    ) + [
        (1, 848, 1264),
        (1, 1264, 848),
        (1, 896, 1200),
        (1, 1200, 896),
        (1, 928, 1152),
        (1, 1152, 928),
        (1, 768, 1376),
        (1, 1376, 768),
        (1, 1584, 672),
    ]

    bucket_group = BucketGroup(bucket_configs)
    img_h, img_w = img.size[1], img.size[0]
    bucket = bucket_group.find_best_bucket((1, 1, img_h, img_w))
    return bucket[-2], bucket[-1]


def resize_center_crop(img: Image.Image, target_size: tuple[int, int]) -> Image.Image:
    width, height = img.size
    target_h, target_w = target_size
    scale = max(target_h / height, target_w / width)
    resize_h, resize_w = math.ceil(height * scale), math.ceil(width * scale)
    img = img.resize((resize_w, resize_h), Image.LANCZOS)
    return TF.center_crop(img, target_size)


def resize_for_bucket(img: Image.Image) -> Image.Image:
    target_h, target_w = get_bucket_size(img)
    return resize_center_crop(img, (target_h, target_w))


def _build_prompt(custom_prompt: str | None, garment_type: str) -> str:
    prompt = (custom_prompt or "").strip()
    if prompt:
        return prompt
    return DEFAULT_PROMPTS.get(garment_type, DEFAULT_PROMPTS["full_body"])


def _run_pipeline(
    model: LoadedTryOnModel,
    processed_images: list[Image.Image],
    prompt_text: str,
    target_h: int,
    target_w: int,
    num_steps: int,
    guidance: float,
    seed: int,
    negative_prompt: str,
) -> Image.Image:
    image_tags = "".join(["<image>\n" for _ in range(len(processed_images))])
    final_prompt = f"<|im_start|>user\n{image_tags}{prompt_text}<|im_end|>\n"
    final_neg = f"<|im_start|>user\n{image_tags}{negative_prompt}<|im_end|>\n"

    result = model.pipeline(
        prompt=[final_prompt],
        negative_prompt=[final_neg],
        images=[processed_images],
        height=[target_h],
        width=[target_w],
        num_frames=1,
        num_inference_steps=int(num_steps),
        guidance_scale=float(guidance),
        generator=torch.Generator(device=model.device).manual_seed(int(seed)),
        output_type="pt",
        return_dict=False,
        enable_denormalization=False,
    )
    video_tensor = result[0]
    img_tensor = (video_tensor[:, 0] * 255).clamp(0, 255).to(torch.uint8).cpu()
    return Image.fromarray(img_tensor.permute(1, 2, 0).numpy())


def _try_crop_and_paste(
    model: LoadedTryOnModel,
    person_image: Image.Image,
    garment_image: Image.Image,
    prompt: str,
    num_inference_steps: int,
    guidance_scale: float,
    seed: int,
    negative_prompt: str,
) -> tuple[Image.Image, str] | None:
    try:
        import cv2
        from xvideo.tryon_utils import create_feathered_mask, get_model_bbox_info
    except Exception:
        return None

    bucket_info = {
        "1.00": [1024, 1024],
        "0.67": [848, 1264],
        "1.49": [1264, 848],
        "0.75": [896, 1200],
        "1.34": [1200, 896],
        "0.81": [928, 1152],
        "1.24": [1152, 928],
        "0.56": [768, 1376],
        "1.79": [1376, 768],
        "2.36": [1584, 672],
    }

    model_bgr = cv2.cvtColor(np.array(person_image), cv2.COLOR_RGB2BGR)
    bbox_success, _original_bbox, extended_bbox, _bucket_size, big_model, _inner_bbox, feather_radius = get_model_bbox_info(
        image=model_bgr,
        bucket_info=bucket_info,
    )
    small_model = person_image.height * person_image.width < 1024**2

    garment_for_model = resize_for_bucket(garment_image)
    if (not bbox_success) or big_model or small_model:
        target_h, target_w = get_bucket_size(person_image)
        input_person = person_image.resize((target_w, target_h))
        output = _run_pipeline(
            model,
            [garment_for_model, input_person],
            prompt,
            target_h,
            target_w,
            num_inference_steps,
            guidance_scale,
            seed,
            negative_prompt,
        )
        return output, "direct path: crop-and-paste skipped because no suitable body bbox was needed/found"

    ex1, ey1, ex2, ey2 = map(int, extended_bbox)
    cropped_bgr = model_bgr[ey1:ey2, ex1:ex2]
    cropped_person = Image.fromarray(cv2.cvtColor(cropped_bgr, cv2.COLOR_BGR2RGB))
    cropped_w, cropped_h = cropped_person.size
    target_h, target_w = get_bucket_size(cropped_person)
    input_person = cropped_person.resize((target_w, target_h))

    output = _run_pipeline(
        model,
        [garment_for_model, input_person],
        prompt,
        target_h,
        target_w,
        num_inference_steps,
        guidance_scale,
        seed,
        negative_prompt,
    )
    resized_output = output.resize((cropped_w, cropped_h))
    mask = create_feathered_mask(
        resized_output.size,
        bbox=extended_bbox,
        feather_radius=feather_radius,
        full_image_size=person_image.size,
    )
    final = person_image.copy()
    final.paste(resized_output, (ex1, ey1), mask)
    return final, "crop-and-paste path: generated cropped body region and pasted it back"


@torch.inference_mode()
def run_tryon(
    person_image: Image.Image,
    garment_images: list[Image.Image],
    garment_type: str = "upper_body",
    primary_garment_index: int = 0,
    seed: int = 42,
    num_inference_steps: int = 30,
    guidance_scale: float = 6.0,
    auto_crop: bool = False,
    prompt: str | None = None,
    negative_prompt: str = DEFAULT_NEG_PROMPT,
    prompt_enhancer: bool = True,
    pe_images: list[Image.Image] | None = None,
) -> tuple[Image.Image, int, str, str, str]:
    if person_image is None:
        raise ValueError("person_image is required.")
    if not garment_images:
        raise ValueError("At least one garment reference image is required.")

    primary_garment, used_index, _all_garments = select_primary_garment(garment_images, primary_garment_index)
    set_seed(seed)
    model = load_model()

    person = person_image.convert("RGB")
    garment = primary_garment.convert("RGB")
    effective_prompt = _build_prompt(prompt, garment_type)
    pe_reference_images = [img.convert("RGB") for img in pe_images] if pe_images else [person, *[img.convert("RGB") for img in garment_images]]
    enhanced_prompt, pe_status = rewrite_prompt_if_enabled(effective_prompt, pe_reference_images, prompt_enhancer)
    effective_negative = (negative_prompt or DEFAULT_NEG_PROMPT).strip()

    if auto_crop:
        cropped_result = _try_crop_and_paste(
            model,
            person,
            garment,
            enhanced_prompt,
            num_inference_steps,
            guidance_scale,
            seed,
            effective_negative,
        )
        if cropped_result is not None:
            image, path_status = cropped_result
            status = f"Success. Used garment index {used_index}. {path_status}."
            return image, used_index, enhanced_prompt, pe_status, status

    target_h, target_w = get_bucket_size(person)
    input_person = resize_center_crop(person, (target_h, target_w))
    input_garment = resize_for_bucket(garment)
    output = _run_pipeline(
        model,
        [input_garment, input_person],
        enhanced_prompt,
        target_h,
        target_w,
        num_inference_steps,
        guidance_scale,
        seed,
        effective_negative,
    )
    status = f"Success. Used garment index {used_index}. direct path."
    return output, used_index, enhanced_prompt, pe_status, status


def clear_model_cache() -> None:
    global _MODEL, _MODEL_DIR
    _MODEL = None
    _MODEL_DIR = None
    if torch.cuda.is_available():
        torch.cuda.empty_cache()