Blame for app.py

"""Universal Lora Studio: Comfy Render Backend + FastAPI Data Backend + Native Static Frontend + ZeroGPU Token Passthrough + Config Driven"""
"""Comfy Version: c1eece62cbf1805d4f64286d7f5592e6a215e6ba"""
import base64
import asyncio
import glob
import json
import os
import pathlib
import random
import re
import shutil
import subprocess
import sys
import traceback
import uuid
import time
import contextvars
import logging
from typing import Any
import requests
from huggingface_hub import hf_hub_download, list_repo_files
from tqdm.auto import tqdm
import tqdm as tqdm_lib
import gradio as gr
from gradio.context import LocalContext
from fastapi import FastAPI, BackgroundTasks, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import uvicorn
import threading

from model_config import MODEL_CONFIG




# ==========================================================================
# 🛠️ Environment Patches & Global Paths 
# ==========================================================================
logging.getLogger("root").setLevel(logging.ERROR) 

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(**_kwargs):
            def decorate(function):
                return function
            return decorate
    spaces = _SpacesFallback()

import torch
if hasattr(torch.cuda, 'get_device_properties'):
    _orig_get_device_properties = torch.cuda.get_device_properties
    def _patched_get_device_properties(device):
        props = _orig_get_device_properties(device)
        if not hasattr(props, 'is_integrated'):
            props.is_integrated = False
        return props
    torch.cuda.get_device_properties = _patched_get_device_properties

ROOT = pathlib.Path(__file__).resolve().parent
COMFY = ROOT / "ComfyUI"
MODELS = COMFY / "models"
INPUT = COMFY / "input"
OUTPUT = COMFY / "output"
COMFY_COMMIT = "c1eece62cbf1805d4f64286d7f5592e6a215e6ba"

_comfy_ready = False
_nodes_ready = False
tasks_status = {}
_last_used_architecture = None
_last_used_unet = None
# _previous_arch_for_gc = None 

# ==========================================================================
# ⚙️ Comfy Config
# ==========================================================================
SAMPLERS = ["euler", "euler_ancestral", "euler_a", "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_sde", "heun", "lms","lcm"]
SCHEDULERS = ["simple", "beta", "normal", "karras", "exponential", "sgm_uniform"]

UNET_CHOICES = {} 
BASE_LORAS_CHOICES = {}
CHAR_LORA_CHOICES = {}

# ==========================================================================
# 📊 Load Config & LoRAs List
# ==========================================================================
def load_config_from_json():
    global UNET_CHOICES, BASE_LORAS_CHOICES
    try:
        print(f"正在获取远端 JSON 配置: {MODEL_CONFIG['config_url']}")
        res = requests.get(MODEL_CONFIG['config_url'], timeout=5)
        if res.ok:
            full_data = res.json()
            root_mapping = MODEL_CONFIG['json_keys']['root_mapping']
            
            for remote_key, local_engine in root_mapping.items():
                engine_data = full_data.get(remote_key, {})
                display_engine = MODEL_CONFIG["models"].get(local_engine, {}).get("name", local_engine)
                
                for item in engine_data.get(MODEL_CONFIG['json_keys']['unets'], []):
                    display_name = f"{display_engine} | {item['name']}"
                    UNET_CHOICES[display_name] = {"repo": item["repo"], "subfolder": item["subfolder"], "filename": item["filename"], "engine": local_engine}
                for item in engine_data.get(MODEL_CONFIG['json_keys']['base_loras'], []):
                    display_name = f"{item['name']}"
                    BASE_LORAS_CHOICES[display_name] = {"repo": item["repo"], "subfolder": item["subfolder"], "filename": item["filename"], "defaultEnable": str(item.get("defaultEnable", "false")).lower() == "true", "engine": local_engine}
            print("✅ 成功加载 JSON 配置")
            return
    except Exception as e:
        print(f"⚠️ JSON 配置加载失败: {e}")

    for engine, cfg in MODEL_CONFIG.get("models", {}).items():
        display_engine = cfg.get("name", engine)
        for name, info in cfg.get("fallback_unets", {}).items():
            UNET_CHOICES[f"{display_engine} | {name}"] = {**info, "engine": engine}
        for name, info in cfg.get("fallback_base_loras", {}).items():
            BASE_LORAS_CHOICES[f"{name}"] = {**info, "engine": engine}

LAST_SCAN_TIME = 0 

def load_char_loras_from_repo():
    global CHAR_LORA_CHOICES, LAST_SCAN_TIME
    current_time = time.time()
    
    if current_time - LAST_SCAN_TIME < 300:
        print("⚡ 使用缓存，跳过重新扫描 LoRA 列表")
        return
        
    print(f"🔄 触发动态刷新，正在重新扫描 {MODEL_CONFIG['char_lora_repo']}...")
    exclude_list = [x.lower() for x in MODEL_CONFIG.get("exclude_folders", [])]
    new_char_choices = {}
    
    try:
        files = list_repo_files(repo_id=MODEL_CONFIG['char_lora_repo'], token=os.environ.get("HF_TOKEN"))
        for engine, cfg in MODEL_CONFIG.get("models", {}).items():
            prefix = cfg.get("char_lora_prefix", "")
            new_char_choices[engine] = {}
            for f in files:
                if f.startswith(prefix) and f.endswith(".safetensors") and "loras/" not in f:
                    relative_path = f[len(prefix):]
                    parts = relative_path.split("/")
                    if len(parts) > 1 and parts[0].lower() in exclude_list:
                        continue 
                    
                    filename = parts[-1]
                    webp_path = f.replace(".safetensors", ".webp")
                    webp_url = f"https://huggingface.co/{MODEL_CONFIG['char_lora_repo']}/resolve/main/{webp_path}"
                    
                    new_char_choices[engine][filename] = {
                        "path": f,
                        "webp": webp_url
                    }
        
        CHAR_LORA_CHOICES.clear()
        CHAR_LORA_CHOICES.update(new_char_choices)
        LAST_SCAN_TIME = current_time
        print(f"✅ 成功刷新角色 LoRA: {[f'{k}: {len(v)}个' for k, v in CHAR_LORA_CHOICES.items()]}")
    except Exception as e:
        print(f"⚠️ 动态刷新失败: {e}")

load_config_from_json()
load_char_loras_from_repo()

# ==========================================================================
# ⬇️ Download Comfy Models & Init
# ==========================================================================
def _run(command: list[str], cwd: pathlib.Path | None = None, check: bool = True) -> None:
    subprocess.run(command, cwd=str(cwd) if cwd else None, check=check)

def _install_filtered_requirements(path: pathlib.Path) -> None:
    if not path.exists(): return
    blocked = {"torch", "torchvision", "torchaudio", "transformers", "huggingface-hub", "accelerate"}
    requirements = []
    for raw in path.read_text(encoding="utf-8", errors="ignore").splitlines():
        item = raw.strip()
        if not item or item.startswith("#"): continue
        package = re.split(r"[<>=!~;\[\s]", item.lower().replace("_", "-"), maxsplit=1)[0]
        if package not in blocked: requirements.append(item)
    if requirements:
        _run([sys.executable, "-m", "pip", "install", "--no-cache-dir", *requirements], check=False)

def _ensure_comfy() -> None:
    global _comfy_ready
    if _comfy_ready: return

    if not COMFY.exists():
        print(f"📦 Cloning ComfyUI repo...")
        _run(["git", "clone", "https://github.com/comfyanonymous/ComfyUI.git", str(COMFY)])
    
    # 强制 Checkout 到指定 Commit
    print(f"🔄 Fetching and checking out ComfyUI commit: {COMFY_COMMIT[:8]}...")
    _run(["git", "fetch", "origin"], cwd=COMFY)
    _run(["git", "checkout", COMFY_COMMIT], cwd=COMFY)

    # 打印当前真实的版本号
    try:
        current_commit = subprocess.check_output(
            ["git", "rev-parse", "HEAD"], cwd=str(COMFY), text=True
        ).strip()
        print(f"✅ ComfyUI 当前版本 Commit: {current_commit}")
    except Exception as e:
        print(f"⚠️ 无法获取 ComfyUI Commit 信息: {e}")

    _install_filtered_requirements(COMFY / "requirements.txt")
    
    for folder in ("diffusion_models", "text_encoders", "vae", "loras"):
        (MODELS / folder).mkdir(parents=True, exist_ok=True)
    INPUT.mkdir(parents=True, exist_ok=True)
    OUTPUT.mkdir(parents=True, exist_ok=True)
    
    _comfy_ready = True

def _init_comfy_nodes() -> None:
    global _nodes_ready
    if _nodes_ready: return
    
    # 1. 获取绝对路径，防止相对路径在某些环境下解析失败
    comfy_path = str(COMFY.resolve())
    
    # 2. 确保绝对路径在 sys.path 的最前面
    if comfy_path not in sys.path:
        sys.path.insert(0, comfy_path)
    
    # 3. 切换工作目录，确保 ComfyUI 内部的相对路径依赖（如 custom_nodes 等）正常
    os.chdir(comfy_path)
    
    # 4. 增加一道硬校验：如果没克隆成功，直接报错提示，而不是抛出难懂的 import error
    if not os.path.exists(os.path.join(comfy_path, "execution.py")):
        raise RuntimeError(f"❌ ComfyUI execution.py not found in {comfy_path}. Did git clone fail?")

    # 5. 开始导入
    import execution, nodes, server
    from app.assets.manager import default_asset_manager
    
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    server_instance = server.PromptServer(loop, default_asset_manager())
    execution.PromptQueue(server_instance)
    loop.run_until_complete(nodes.init_extra_nodes())
    
    _nodes_ready = True
    print("✅ ComfyUI 节点初始化完成！", flush=True)

# ==========================================================================
# ⚡ Comfy optimization
# ==========================================================================
    
# --------------------------------------------------------------------------
# 🩹 Comfy Speedup Patch V3: Aggressive Interception + Precise Fix for INT8 CUDA Tensor Error
# Commit @11731865e8832a9b2edffb05462621223097e908
# --------------------------------------------------------------------------
def patch_comfy_loader():
    print("[🚀 提速补丁] 开始注入...", flush=True)
    try:
        comfy_path = str(COMFY.resolve())
        if comfy_path not in sys.path:
            sys.path.insert(0, comfy_path)
            
        print("[🚀 提速补丁] [1/3] 正在导入 comfy.utils 等基础模块...", flush=True)
        import comfy.utils
        import safetensors.torch
        import safetensors
        import comfy.ldm.modules.attention
        from comfy.ldm.modules.attention import OptimizedAttention
        print("[🚀 提速补丁] [1/3] ✅ 基础模块导入成功", flush=True)
        
        _original_load_torch_file = comfy.utils.load_torch_file
        _original_st_load_file = safetensors.torch.load_file
        
        def _fast_load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
            if device is None:
                import comfy.model_management
                device = comfy.model_management.get_torch_device()
                
            if not safe_load:
                return _original_load_torch_file(ckpt, safe_load, device, return_metadata=return_metadata)
                
            file_path = pathlib.Path(ckpt)
            if not file_path.exists():
                return _original_load_torch_file(ckpt, safe_load, device, return_metadata=return_metadata)

            if file_path.suffix == ".safetensors" and file_path.stat().st_size > 500 * 1024 * 1024:
                print(f"[🚀 提速补丁] 暴力读取大模型: {file_path.name} ({file_path.stat().st_size / 1024 / 1024:.1f} MB)", flush=True)
                try:
                    with open(file_path, "rb") as f:
                        data = f.read()
                    
                    header_size = int.from_bytes(data[:8], byteorder="little")
                    header = json.loads(data[8:8+header_size].decode("utf-8"))
                    
                    # 正确提取 header 里的 __metadata__，这是合法的 JSON 字符串字典
                    metadata = header.pop("__metadata__", {})
                    
                    tensors = {}
                    data_offset = 8 + header_size
                    
                    for tensor_name, info in header.items():
                        # 🌟 关键修复：如果权重名恰好叫 metadata，直接跳过，绝不把它当权重读进显存
                        if tensor_name in ("metadata", "__metadata__"):
                            continue

                        dtype = info["dtype"]
                        shape = info["shape"]
                        start, end = info["data_offsets"]
                        
                        tensor_bytes = data[data_offset + start : data_offset + end]

                        dtype_map = {
                            "F16": torch.float16, "BF16": torch.bfloat16,
                            "F32": torch.float32, "I8": torch.int8,
                            "U8": torch.uint8, "I32": torch.int32
                        }
                        torch_dtype = dtype_map.get(dtype, torch.float16)

                        # 🌟 P1: 大 tensor 走 numpy 零拷贝，避免 bytearray 二次复制
                        if len(tensor_bytes) > 50 * 1024 * 1024:
                            import numpy as np
                            np_dtype_map = {
                                "F16": np.float16, "BF16": np.uint16, "F32": np.float32,
                                "I8": np.int8, "U8": np.uint8, "I32": np.int32,
                            }
                            np_dtype = np_dtype_map.get(dtype)
                            if np_dtype is not None:
                                np_arr = np.frombuffer(tensor_bytes, dtype=np_dtype)
                                t = torch.from_numpy(np_arr)
                                if dtype == "BF16":
                                    t = t.view(torch.bfloat16)
                                # clone 强制拥有独立内存，切断与 data 的共享，防止 GC 后 use-after-free
                                tensors[tensor_name] = t.reshape(shape).clone()
                            else:
                                tensors[tensor_name] = torch.frombuffer(bytearray(tensor_bytes), dtype=torch_dtype).reshape(shape)
                        else:
                            tensors[tensor_name] = torch.frombuffer(bytearray(tensor_bytes), dtype=torch_dtype).reshape(shape)
                        
                    print(f"[🚀 提速补丁] 暴力读取完成: {file_path.name}", flush=True)
                    
                    if return_metadata:
                        return tensors, metadata
                    return tensors
                    
                except Exception as e:
                    print(f"[⚠️ 提速补丁] 暴力读取失败，回退默认方式: {e}", flush=True)
                    return _original_load_torch_file(ckpt, safe_load, device, return_metadata=return_metadata)
            
            return _original_load_torch_file(ckpt, safe_load, device, return_metadata=return_metadata)

        print("[🚀 提速补丁] [2/3] 正在替换 comfy.utils.load_torch_file...", flush=True)
        comfy.utils.load_torch_file = _fast_load_torch_file
        print("[🚀 提速补丁] [2/3] ✅ load_torch_file 已替换", flush=True)
        
        def _wrapped_st_load_file(filename, device='cpu'):
            if str(filename).endswith(".safetensors") and os.path.exists(filename) and os.path.getsize(filename) > 500 * 1024 * 1024:
                return _fast_load_torch_file(filename, safe_load=True, device=device, return_metadata=False)
            return _original_st_load_file(filename, device)
            
        safetensors.torch.load_file = _wrapped_st_load_file
        print("[🚀 提速补丁] [2/3] ✅ safetensors.torch.load_file 已替换", flush=True)

        # =========================================================
        # 🌟 V3 Core Fix: More aggressive interception of OptimizedAttention._load_from_state_dict
        # =========================================================
        print("[🚀 提速补丁] [3/3] 正在替换 OptimizedAttention._load_from_state_dict...", flush=True)
        _original_load_from_state_dict = OptimizedAttention._load_from_state_dict

        def _patched_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
            metadata_key = prefix + "metadata"
            
            # 如果你有暴力读取留下的 metadata，或者有其他任何叫 metadata 的 Tensor
            if metadata_key in state_dict:
                metadata = state_dict[metadata_key]
                try:
                    # 强行拉回 CPU 转 Numpy 提取 JSON
                    if hasattr(metadata, "cpu"):
                        metadata_cpu = metadata.cpu()
                        if hasattr(metadata_cpu, "numpy"):
                            self.config = json.loads(metadata_cpu.numpy().tobytes())
                        else:
                            self.config = json.loads(metadata_cpu.tobytes())
                    elif isinstance(metadata, bytes) or isinstance(metadata, str):
                        self.config = json.loads(metadata)
                except Exception as e:
                    print(f"[🚀 提速补丁] 解析 attention config 失败: {e}", flush=True)
                
                # 处理完立刻删掉，防止底层报 Unexpected key 或者重复处理报错
                del state_dict[metadata_key]
            
            # 如果 local_metadata 里有（有些版本从这里拿），也顺手处理一下
            if local_metadata and isinstance(local_metadata, dict) and "__metadata__" in local_metadata:
                try:
                    self.config = json.loads(local_metadata["__metadata__"])
                except:
                    pass

            # 交还给原本的逻辑（此时 state_dict 已经干净了）
            return _original_load_from_state_dict(
                self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
            )

        OptimizedAttention._load_from_state_dict = _patched_load_from_state_dict
        print("[🚀 提速补丁] [3/3] ✅ OptimizedAttention 防御已替换", flush=True)
        # =========================================================

        # =========================================================
        # 🌟 P1 真探针：patch model_management.load_models_gpu 测真实 load 耗时
        # =========================================================
        import comfy.model_management as mm
        _original_load_models_gpu = mm.load_models_gpu

        def _timed_load_models_gpu(models, *args, **kwargs):
            t0 = time.time()
            result = _original_load_models_gpu(models, *args, **kwargs)
            elapsed = round(time.time() - t0, 2)
            task_id = _current_generation_task_id.get()
            if task_id:
                _write_metric(task_id, load_time=elapsed, load_time_source="load_models_gpu")
                print(f"[探针] 🎯 load_models_gpu 耗时: {elapsed}s", flush=True)
            return result

        mm.load_models_gpu = _timed_load_models_gpu
        print("[🚀 提速补丁] [探针] ✅ load_models_gpu 已接管", flush=True)
        # =========================================================

        print("✅ [🚀 提速补丁] V3 部署完成！INT8 Tensor 防御已就绪！", flush=True)
    except Exception as e:
        # 👇 下面只增加了解析报错原因和状态提示的 print，没有修改任何原有逻辑
        if "cannot import name 'OptimizedAttention'" in str(e):
            print(f"⚠️ [🚀 提速补丁] [3/3] 跳过: 当前 ComfyUI 版本未提供 OptimizedAttention，防御补丁未启用", flush=True)
            print(f"✅ [🚀 提速补丁] 核心提速已生效！load_torch_file 快路径接管成功，不影响正常生成！", flush=True)
        else:
            print(f"⚠️ [🚀 提速补丁] 注入失败: {e}", flush=True)

def _download_if_needed(repo: str, subfolder: str, filename: str, dest_root: pathlib.Path) -> None:
    dest_path = dest_root / filename
    if dest_path.exists(): return
    print(f"⏳ Downloading: {filename} ...")
    dest_root.mkdir(parents=True, exist_ok=True)
    hf_hub_download(repo_id=repo, subfolder=subfolder if subfolder else None, filename=filename, local_dir=str(dest_root))
    src = dest_root / subfolder / filename if subfolder else dest_root / filename
    if src.exists() and src != dest_path:
        shutil.move(str(src), str(dest_path))

def _ensure_models() -> None:
    print("开始全量下载基础模型...")
    for engine_cfg in MODEL_CONFIG["models"].values():
        for repo, file_path in engine_cfg.get("companion_downloads", []):
            path = pathlib.Path(file_path)
            if "text_encoders" in str(path): dest_root = MODELS / "text_encoders"
            elif "vae" in str(path): dest_root = MODELS / "vae"
            else: dest_root = MODELS / path.parent
            _download_if_needed(repo, str(path.parent), path.name, dest_root)
    for name, cfg in UNET_CHOICES.items():
        _download_if_needed(cfg["repo"], cfg["subfolder"], cfg["filename"], MODELS / "diffusion_models")
    lora_root = MODELS / "loras"
    for name, cfg in BASE_LORAS_CHOICES.items():
        _download_if_needed(cfg["repo"], cfg["subfolder"], cfg["filename"], lora_root)
    print("✅ 基础模型下载完毕！")

# --------------------------------------------------------------------------
# ♻️  Reusable Executor Singleton (preserves Comfy model cache across tasks)
# --------------------------------------------------------------------------
import threading

_exec_lock = threading.Lock()          # 串行化执行，GPU 一次一个任务
_exec_state = {
    "loop": None,
    "server": None,
    "executor": None,
}

_init_lock = threading.Lock()

def _get_executor():
    if _exec_state["executor"] is not None:
        return _exec_state["loop"], _exec_state["executor"]
    with _init_lock:
        if _exec_state["executor"] is not None:
            return _exec_state["loop"], _exec_state["executor"]

        # 👇 从这里开始全部缩进 4 格
        if str(COMFY) not in sys.path:
            sys.path.insert(0, str(COMFY))
        import execution, server
        from app.assets.manager import default_asset_manager

        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        server_instance = server.PromptServer(loop, default_asset_manager())

        executor = execution.PromptExecutor(
            server_instance,
            cache_type=execution.CacheType.RAM_PRESSURE,
            cache_args={"lru": 8, "ram": 16.0, "ram_inactive": 0},
        )

        _exec_state.update(loop=loop, server=server_instance, executor=executor)
        print("[执行器] ♻️ 单例执行器已创建（将跨任务复用）", flush=True)
        # 👇 新增这一行，方便你在日志里观察当前 VRAM 状态
        try:
            import torch
            allocated = torch.cuda.memory_allocated() / (1024**3)
            reserved = torch.cuda.memory_reserved() / (1024**3)
            print(f"[显存监控] 已分配: {allocated:.2f} GB | 已保留: {reserved:.2f} GB", flush=True)
        except Exception:
            pass
    return loop, executor
# --------------------------------------------------------------------------
# 🧩  Loader Chain Cache (stable node ids => Comfy node-cache hits, no re-patch
# --------------------------------------------------------------------------

_loader_chain_cache: dict[str, tuple[dict, str]] = {}

def _combo_key(engine: str, unet_name: str,
               base_loras: list[tuple[str, float]], char_lora: str | None, char_strength: float) -> str:
    base = tuple((f, round(float(s), 3)) for f, s in base_loras)
    return f"{engine}|{unet_name}|{base}|{char_lora or ''}|{round(float(char_strength), 3)}"

def _build_loader_chain(engine, unet_name, base_loras, char_lora, char_strength):
    """
    只构建 1..N 的 loader 链（unet/clip/vae + lora 叠加），
    返回 (chain_nodes, final_model_node_id)。
    chain_nodes 的节点 id 从 "1" 开始连续，保证相同组合 -> 相同 id/inputs。
    """
    engine_cfg = MODEL_CONFIG["models"].get(engine, {})
    nodes = MODEL_CONFIG["node_types"]
    clip_loader_node = engine_cfg.get("clip_loader_override", nodes["clip_loader"])
    vae_filename = engine_cfg.get("local_vae_filename", "")

    chain: dict[str, Any] = {
        "1": {"class_type": nodes["unet_loader"], "inputs": {"unet_name": unet_name, "weight_dtype": "default"}},
        "2": {"class_type": clip_loader_node, "inputs": _build_clip_inputs(engine_cfg, clip_loader_node)},
        "3": {"class_type": nodes["vae_loader"], "inputs": {"vae_name": vae_filename}},
    }
    cur = "1"
    idx = 4
    for lora_file, strength in base_loras:
        chain[str(idx)] = {"class_type": nodes["lora_loader"],
                           "inputs": {"model": _ref(cur), "lora_name": lora_file, "strength_model": strength}}
        cur = str(idx); idx += 1
    if char_lora:
        chain[str(idx)] = {"class_type": nodes["lora_loader"],
                           "inputs": {"model": _ref(cur), "lora_name": char_lora, "strength_model": char_strength}}
        cur = str(idx); idx += 1
    return chain, cur

# ==========================================================================
# 💻 Prepare runtime
# ==========================================================================

def _warmup_node_imports() -> None:
    """预热 Comfy 节点类导入，消除首次 execute 的 import 开销"""
    try:
        import nodes as comfy_nodes
        for class_name in ["UNETLoader", "CLIPLoader", "DualCLIPLoader", "VAELoader",
                           "LoraLoaderModelOnly", "KSampler", "VAEDecode", "SaveImage",
                           "CLIPTextEncode", "ConditioningZeroOut", "EmptyLatentImage"]:
            if hasattr(comfy_nodes, class_name):
                getattr(comfy_nodes, class_name)
        print("[预热] ✅ 节点类导入完成", flush=True)
    except Exception as e:
        print(f"[预热] ⚠️ {e}", flush=True)


def _prepare_runtime() -> None:
    _ensure_comfy()
    _ensure_models()
    _init_comfy_nodes()
    patch_comfy_loader()
    _warmup_node_imports()          # 🌟 P1: 预热节点类

try:
    print("⬇️ Starting up...")
    _prepare_runtime()
except Exception as e:
    print(f"⚠️ Startup failed: {e}")


# ==========================================================================
# 🔄 Progress Interceptor
# ==========================================================================
_current_generation_task_id = contextvars.ContextVar("current_generation_task_id", default=None)
_current_download_task_id = contextvars.ContextVar("current_download_task_id", default=None)

class UnifiedProgressTqdm(tqdm):
    def update(self, n=1):
        super().update(n)

        download_task_id = _current_download_task_id.get()
        gen_task_id = _current_generation_task_id.get()

        if download_task_id and download_task_id in tasks_status:
            if self.total and self.total > 0:
                percent = int((self.n / self.total) * 100)
                tasks_status[download_task_id]["message"] = f"📥 Downloading LoRA: {percent}%..."
            else:
                tasks_status[download_task_id]["message"] = f"📥 Downloading LoRA: {self.n} units..."
            return 

        if gen_task_id and gen_task_id in tasks_status:
            try:
                metric_file = f"/tmp/metrics_{gen_task_id}.json"
                if os.path.exists(metric_file):
                    with open(metric_file, "r") as f:
                        saved_data = json.load(f)

                    # 🌟 首次 update 就落 load_time，但优先保留 load_models_gpu 的真值
                    if "load_time" not in saved_data or saved_data.get("load_time", -1) == -1:
                        # 如果 load_models_gpu 还没触发，这里用 tqdm 的首次 update 做兜底
                        if saved_data.get("load_time_source") != "load_models_gpu":
                            load_start = saved_data.get("load_start", time.time())
                            saved_data["load_time"] = round(time.time() - load_start, 2)
                            saved_data["load_time_source"] = "tqdm_first_update"
                            print(f"[探针] 🎯 tqdm 首次 update load_time: {saved_data['load_time']}s", flush=True)

                    if self.total and self.total > 1:
                        percent = int((self.n / self.total) * 100)
                        saved_data["phase"] = f"🎨 Rendering: {self.n}/{self.total} steps ({percent}%)"

                    with open(metric_file, "w") as f:
                        json.dump(saved_data, f)
            except Exception as e:
                print(f"[探针] ❌ 更新物理文件失败: {e}", flush=True)

            if self.total and self.total > 1:
                percent = int((self.n / self.total) * 100)
                current_step_msg = f"🎨 Rendering: {self.n}/{self.total} steps ({percent}%)"
                tasks_status[gen_task_id]["message"] = current_step_msg
                print(f"[进度拦截] Task {gen_task_id[:8]} | {self.n}/{self.total} steps ({percent}%)")

# ==========================================================================
# ⚛️ Comfy Backend & Workflow
# ==========================================================================
                
import psutil  

def print_ram_status(stage: str):
    mem = psutil.virtual_memory()
    process = psutil.Process(os.getpid())
    proc_mem_gb = process.memory_info().rss / (1024 ** 3)
    total_gb = mem.total / (1024 ** 3)
    avail_gb = mem.available / (1024 ** 3)
    
    # 🌟 新增：读取真实的 Cgroup 限制
    cgroup_limit_gb = -1
    try:
        # 尝试读取 Cgroup v2
        with open("/sys/fs/cgroup/memory.max", "r") as f:
            val = f.read().strip()
            if val != "max":
                cgroup_limit_gb = int(val) / (1024 ** 3)
    except FileNotFoundError:
        try:
            # 尝试读取 Cgroup v1
            with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "r") as f:
                val = f.read().strip()
                # v1 经常用一个巨大的数字表示无限制
                if int(val) < (1 << 60):
                    cgroup_limit_gb = int(val) / (1024 ** 3)
        except Exception:
            pass

    cgroup_str = f" | 容器真实限额: {cgroup_limit_gb:.2f} GB" if cgroup_limit_gb > 0 else " | 容器真实限额: 未知/无限制"
    print(f"[RAM监控 - {stage}] 宿主机总内存: {total_gb:.2f} GB | 宿主机可用: {avail_gb:.2f} GB | 本进程占用: {proc_mem_gb:.2f} GB | 系统已用: {mem.percent}%{cgroup_str}")


def _write_metric(task_id: str, **updates):
    if not task_id:
        return
    try:
        metric_file = f"/tmp/metrics_{task_id}.json"
        saved_data = {}
        if os.path.exists(metric_file):
            with open(metric_file, "r") as f:
                saved_data = json.load(f)
        saved_data.update(updates)
        with open(metric_file, "w") as f:
            json.dump(saved_data, f)
    except Exception as e:
        print(f"[探针] 写入 metric 失败: {e}", flush=True)

def _ref(node: str, output: int = 0) -> list[Any]:
    return [node, output]

def _log_workflow_summary(workflow: dict[str, Any], task_id: str, load_time_val: float, wf_time: float):
    """独立提取并打印单行 JSON 摘要日志（不影响主流程逻辑）"""
    try:
        # 获取当前日期时间
        current_time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        
        # 1. 提取 UNet
        unet_node = next((v for v in workflow.values() if v.get("class_type") in ["UNETLoader", "CheckpointLoaderSimple"]), None)
        unet_name = unet_node["inputs"].get("unet_name", unet_node["inputs"].get("ckpt_name", "Unknown")) if unet_node else "Unknown"
        
        # 2. 提取 VAE
        vae_node = next((v for v in workflow.values() if v.get("class_type") == "VAELoader"), None)
        vae_name = vae_node["inputs"].get("vae_name", "Unknown") if vae_node else "Unknown"

        # 3. 提取 CLIP / TE
        te_name = "Unknown"
        clip_node = next((v for v in workflow.values() if v.get("class_type") in ["CLIPLoader", "DualCLIPLoader"]), None)
        if clip_node:
            if clip_node["class_type"] == "DualCLIPLoader":
                te_name = f"{clip_node['inputs'].get('clip_name1')} + {clip_node['inputs'].get('clip_name2')}"
            else:
                te_name = clip_node["inputs"].get("clip_name", "Unknown")

        # 4. 提取 LoRA
        loras = []
        for v in workflow.values():
            if v.get("class_type") in ["LoraLoader", "LoraLoaderModelOnly"]:
                l_name = v["inputs"].get("lora_name", "Unknown").replace(".safetensors", "")
                l_str = v["inputs"].get("strength_model", 1.0)
                loras.append(f"{l_name}({l_str})")
        lora_str = ",".join(loras) if loras else "None"
        
        # 5. 提取 Sampler 参数和尺寸
        sampler_node = next((v for v in workflow.values() if v.get("class_type") == "KSampler"), None)
        latent_node = next((v for v in workflow.values() if v.get("class_type") == "EmptyLatentImage"), None)
        
        sampler_str, steps, cfg, seed, w, h = "", 0, 0, 0, 0, 0
        if sampler_node:
            s = sampler_node["inputs"]
            sampler_str = f"{s.get('sampler_name')}/{s.get('scheduler')}"
            steps, cfg, seed = s.get('steps'), s.get('cfg'), s.get('seed')
        if latent_node:
            w, h = latent_node["inputs"].get("width"), latent_node["inputs"].get("height")

        # 6. 提取 Prompt (不截断，完全显示)
        prompt_text = "None"
        for v in workflow.values():
            if v.get("class_type") == "CLIPTextEncode":
                p = v["inputs"].get("text", "")
                prompt_text = p
                break
        # 去除换行符，防止破坏单行展示；如果有单引号，转义一下
        prompt_text = prompt_text.replace("\n", " ").replace("'", "\\'")

        task_short = task_id[:8] if task_id else 'None'
        
        # 打印单行 JSON
        print(f"✅ [{task_short}] {{"
              f"'date': '{current_time_str}', "
              f"'unet': '{unet_name}', "
              f"'te': '{te_name}', "
              f"'vae': '{vae_name}', "
              f"'lora': '{lora_str}', "
              f"'size': '{w}x{h}', "
              f"'seed': {seed}, "
              f"'sampler': '{sampler_str} {steps}s cfg{cfg}', "
              f"'prompt': '{prompt_text}', "
              f"'load': {load_time_val}s, "
              f"'total': {wf_time}s"
              f"}}", flush=True)
              
    except Exception as e:
        print(f"[Log Format Error]: {e}", flush=True)


def _execute_workflow(workflow: dict[str, Any], task_id: str = None) -> list[str]:
    print_ram_status("任务开始前")
    import tqdm as tqdm_lib
    import tqdm.auto
    
    original_tqdm = tqdm_lib.tqdm
    original_auto_tqdm = tqdm.auto.tqdm
    
    tqdm_lib.tqdm = UnifiedProgressTqdm
    tqdm.auto.tqdm = UnifiedProgressTqdm
    
    try:
        loop, executor = _get_executor()
        prompt_id = str(uuid.uuid4())
        save_id = next(k for k, v in workflow.items() if v.get("class_type") == "SaveImage")
        
        if task_id and task_id in tasks_status:
            tasks_status[task_id]["message"] = "🔄 Executing Workflow..."
            try:
                metric_file = f"/tmp/metrics_{task_id}.json"
                saved_data = {}
                if os.path.exists(metric_file):
                    with open(metric_file, "r") as f:
                        saved_data = json.load(f)
                saved_data["phase"] = "🔄 Executing Workflow..." 
                saved_data["load_start"] = time.time() 
                with open(metric_file, "w") as f:
                    json.dump(saved_data, f)
            except Exception:
                pass
        

        with _exec_lock:                      # ← 串行，复用同一 executor
            # 清理上一次的 history/status，避免状态串味
            executor.history_result = {}
            executor.status_messages = []
            executor.success = True

            workflow_start = time.time()

            executor.execute(workflow, prompt_id, extra_data={}, execute_outputs=[save_id])
            wf_time = round(time.time() - workflow_start, 2)

        _write_metric(task_id, workflow=wf_time)
        
        # 兜底补录 load_time 逻辑保持不变
        load_time_val = 0
        if task_id:
            try:
                metric_file = f"/tmp/metrics_{task_id}.json"
                saved_data = {}
                if os.path.exists(metric_file):
                    with open(metric_file, "r") as f:
                        saved_data = json.load(f)
                if "load_time" not in saved_data or saved_data.get("load_time", -1) == -1:
                    # 真正的 load 耗时由 load_models_gpu 探针写入；这里只标记缺失
                    saved_data["load_time"] = -1
                    saved_data["load_time_source"] = "cache_hit"
                    print(f"[探针] ⚡ 缓存命中（或无需重新加载显存），耗时 < 1s", flush=True)
                    with open(metric_file, "w") as f:
                        json.dump(saved_data, f)
                load_time_val = saved_data.get("load_time", 0)
            except Exception as e:
                print(f"写入 workflow 指标失败: {e}")

        _log_workflow_summary(workflow, task_id, load_time_val, wf_time)

        print_ram_status("工作流执行完毕后")

        if not executor.success:
            message = executor.status_messages[-1] if executor.status_messages else "ComfyUI execution failed"
            raise RuntimeError(str(message))
            
        paths = []
        for output in executor.history_result.get("outputs", {}).values():
            for items in output.values():
                if not isinstance(items, list): continue
                for item in items:
                    if not isinstance(item, dict) or not item.get("filename"): continue
                    candidate = OUTPUT / item.get("subfolder", "") / item["filename"]
                    if candidate.exists(): paths.append(str(candidate))
        if not paths:
            paths = sorted([p for p in glob.glob(str(OUTPUT / "**" / "*.png"), recursive=True)], key=lambda p: os.path.getmtime(p), reverse=True)
        return paths[:1]
    finally:
        tqdm_lib.tqdm = original_tqdm
        tqdm.auto.tqdm = original_auto_tqdm

def prepare_env(char_enable: bool, char_lora: str, unet_choice: str, task_id: str = None):
    _current_download_task_id.set(task_id)
    
    import tqdm.auto
    original_tqdm = tqdm.auto.tqdm
    tqdm.auto.tqdm = UnifiedProgressTqdm
    
    try:
        engine = UNET_CHOICES.get(unet_choice, {}).get("engine", "MOSS")
        if char_enable and char_lora:
            lora_root = MODELS / "loras"
            dest = lora_root / char_lora
            if not dest.exists():
                file_path = CHAR_LORA_CHOICES.get(engine, {}).get(char_lora, {}).get("path")
                if file_path:
                    if task_id and task_id in tasks_status:
                        tasks_status[task_id]["message"] = f"📥 Downloading LoRA ({char_lora})..."
                    
                    hf_hub_download(repo_id=MODEL_CONFIG['char_lora_repo'], filename=file_path, local_dir=str(lora_root))
                    src = lora_root / file_path
                    if src.exists(): shutil.move(str(src), str(dest))
                    
                    if task_id and task_id in tasks_status:
                        tasks_status[task_id]["message"] = "📦 LoRA downloaded, loading..."
    finally:
        tqdm.auto.tqdm = original_tqdm
        _current_download_task_id.set(None)

def select_base_loras(engine: str, lora_states: tuple) -> list[tuple[str, float]]:
    selected = []
    
    # 找出当前引擎专属的 Base LoRA 列表，保持顺序
    valid_loras = []
    for name, cfg in BASE_LORAS_CHOICES.items():
        if cfg["engine"] == engine:
            valid_loras.append((name, cfg))
            
    # 解析前端传来的勾选状态
    for i in range(0, len(lora_states), 2):
        enable = lora_states[i]
        try:
            strength = float(lora_states[i + 1])
        except (ValueError, TypeError):
            strength = 1.0
            
        if enable and strength > 0:
            # 如果前端的数组长度和当前引擎实际可用的 LoRA 列表长度一致
            idx = i // 2
            if idx < len(valid_loras):
                lora_name, lora_cfg = valid_loras[idx]
                selected.append((lora_cfg["filename"], strength))
                
                # 你的探针（保留观察）
                print(f"[LoRA选择探针] ✅ 成功匹配并加载: {lora_name} ({strength})")
                
    return selected

def _build_clip_inputs(engine_cfg: dict, clip_loader_node: str) -> dict:
    clip_type = engine_cfg.get("clip_type", "moss_image") 
    clip_inputs = {"type": clip_type}
    te_filenames = engine_cfg.get("local_te_filename", [])
    if clip_loader_node == "DualCLIPLoader":
        if len(te_filenames) >= 2:
            if "clip_l" in te_filenames[0].lower():
                clip_inputs["clip_name1"] = te_filenames[0]
                clip_inputs["clip_name2"] = te_filenames[1]
            else:
                clip_inputs["clip_name1"] = te_filenames[1]
                clip_inputs["clip_name2"] = te_filenames[0]
        elif len(te_filenames) == 1:
            clip_inputs["clip_name1"] = te_filenames[0]
    else:
        if te_filenames: clip_inputs["clip_name"] = te_filenames[0]
    if "clip_device" in MODEL_CONFIG.get("loader_args", {}):
            clip_inputs["device"] = MODEL_CONFIG["loader_args"]["clip_device"]
    return clip_inputs

def build_workflow(
    prompt: str,
    unet_name: str,
    engine: str,
    width: int,
    height: int,
    steps: int,
    cfg: float,
    seed: int,
    sampler: str,
    scheduler: str,
    base_loras: list[tuple[str, float]],
    char_lora: str | None,
    char_strength: float,
) -> dict[str, Any]:
    # --- loader 链：命中缓存则复用（节点 id / inputs 完全一致） ---
    key = _combo_key(engine, unet_name, base_loras, char_lora, char_strength)
    cached = _loader_chain_cache.get(key)
    if cached is not None:
        chain, current_model_node = cached
        workflow = dict(chain)  # 浅拷贝，下游不改 loader 链
    else:
        chain, current_model_node = _build_loader_chain(
            engine, unet_name, base_loras, char_lora, char_strength)
        _loader_chain_cache[key] = (chain, current_model_node)
        # 只保留最近 8 个组合，避免内存无限增长
        if len(_loader_chain_cache) > 8:
            _loader_chain_cache.pop(next(iter(_loader_chain_cache)))
        workflow = dict(chain)

    # --- 下游节点：id 从链末端续接，与原逻辑完全一致 ---
    nodes = MODEL_CONFIG["node_types"]
    node_idx = max(int(k) for k in chain.keys()) + 1

    final_prompt = prompt.strip()
    workflow[str(node_idx)] = {"class_type": "CLIPTextEncode", "inputs": {"clip": _ref("2"), "text": final_prompt}}
    workflow[str(node_idx+1)] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": _ref(str(node_idx))}}
    workflow[str(node_idx+2)] = {"class_type": "EmptyLatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}}

    workflow[str(node_idx+3)] = {"class_type": "KSampler", "inputs": {"model": _ref(current_model_node), "positive": _ref(str(node_idx)), "negative": _ref(str(node_idx+1)), "latent_image": _ref(str(node_idx+2)), "seed": seed, "steps": steps, "cfg": cfg, "sampler_name": sampler, "scheduler": scheduler, "denoise": 1.0}}
    workflow[str(node_idx+4)] = {"class_type": "VAEDecode", "inputs": {"samples": _ref(str(node_idx+3)), "vae": _ref("3")}}
    workflow[str(node_idx+5)] = {"class_type": "SaveImage", "inputs": {"images": _ref(str(node_idx+4)), "filename_prefix": MODEL_CONFIG["name"].lower()}}
    return workflow

# 🌟 增加全局变量用于追踪 GPU 域内的架构切换

@spaces.GPU(duration=40)  # 🌟 降低 duration 到 40
def generate(
    prompt: str, unet_choice: str, custom_w: int, custom_h: int,
    steps: int, cfg: float, seed: int, sampler: str, scheduler: str,
    char_enable: bool, char_lora: str, char_strength: float,
    *lora_states,
    task_id: str = None,
) -> str:

    # global _previous_arch_for_gc
    
    print(f"[Token探针] task={task_id} | token_present={bool(LocalContext.request.get(None) and LocalContext.request.get(None).headers.get('x-ip-token'))}", flush=True)

    # 获取当前架构
    # current_arch = UNET_CHOICES.get(unet_choice, {}).get("engine", "unknown")

    # 🚫 【已注释】原因：经过极限测试，跨架构排队时会导致 ZeroGPU 重新申请显存，拖慢加载。
    # 交给 ComfyUI 原生的 RAM_PRESSURE 机制（lru=8, ram=16.0）平滑过渡效果更好。
    # if _previous_arch_for_gc is not None and _previous_arch_for_gc != current_arch:
    #     import gc
    #     if torch.cuda.is_available():
    #         torch.cuda.empty_cache()
    #         gc.collect()
    #         print(f"[显存清理] 架构切换 {_previous_arch_for_gc} -> {current_arch}，已清理碎片", flush=True)
    #
    # _previous_arch_for_gc = current_arch

    if task_id:
        if task_id not in tasks_status:
            tasks_status[task_id] = {}
            
        now_ts = time.time()
        start_ts = tasks_status[task_id].get("queue_start_time", now_ts)
        raw_diff = now_ts - start_ts
        
        print(f"[排队探针] 拿到卡时间: {now_ts:.2f}, 排队起点时间: {start_ts:.2f}, 原始差值: {raw_diff:.2f}s", flush=True)
        
        q_time = max(0.0, round(raw_diff, 2))
        
        tasks_status[task_id]["metrics_queue"] = q_time
        tasks_status[task_id]["message"] = "⚙️ Allocating GPU..."
        
        try:
            saved_data = {}
            metric_file = f"/tmp/metrics_{task_id}.json"
            if os.path.exists(metric_file):
                with open(metric_file, "r") as f:
                    saved_data = json.load(f)
            
            saved_data["queue"] = q_time
            saved_data["phase"] = "Allocating GPU..." 
            with open(metric_file, "w") as f:
                json.dump(saved_data, f)
        except Exception as e:
            print(f"写入 queue 指标失败: {e}")
            
        print(f"[GPU 探针] 🟢 成功获取 GPU！排队耗时: {q_time}s", flush=True)

    effective_seed = int(seed)
    unet_info = UNET_CHOICES[unet_choice]
    unet_name = unet_info["filename"]
    engine = unet_info["engine"] 
    
    try:

        workflow = build_workflow(
            prompt=prompt, unet_name=unet_name, engine=engine, width=int(custom_w), height=int(custom_h),
            steps=int(steps), cfg=float(cfg), seed=effective_seed, sampler=sampler, scheduler=scheduler,
            base_loras=select_base_loras(engine, lora_states),
            char_lora=char_lora if (char_enable and char_lora) else None,
            char_strength=float(char_strength),
        )

        print(f"[GPU 探针] 📦 模型节点构建完毕，准备调用 ComfyUI 执行引擎...", flush=True)
        img_path = _execute_workflow(workflow, task_id)[0]
        return img_path, workflow
    except Exception as exc:
        print(traceback.format_exc(), flush=True)
        raise gr.Error(f"generation failed: {str(exc)[:500]}") from exc

# ==========================================================================
# 🖥️ Gradio UI (Dummy for ZeroGPU)
# ==========================================================================
with gr.Blocks() as demo:
    pass
demo.queue()

# ==========================================================================
# 🌐 FastAPI & API Routes
# ==========================================================================
class GenerateRequest(BaseModel):
    prompt: str
    unet_choice: str
    custom_w: int = 1024
    custom_h: int = 1024
    steps: int = 8
    cfg: float = 1.0
    seed: int = -1
    sampler: str = "euler"
    scheduler: str = "simple"
    char_enable: bool = True
    char_lora: str = ""
    char_strength: float = 1.0
    base_lora_states: list = []

import threading
# 专门用于保护 tasks_status 字典的状态锁
_status_lock = threading.RLock()
_model_switch_lock = threading.Lock()

def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str = None, request: Request = None):
    global _last_used_architecture, _last_used_unet

    _current_generation_task_id.set(task_id)
    
    # 初始状态
    with _status_lock:
        tasks_status[task_id] = {"status": "processing", "message": "⏳ Queue for GPU..."}
        
    print(f"[Token探针] task={task_id[:8]} | route_token_present={bool(ip_token)}", flush=True)

    try:
        # --- 阶段 1：环境准备与模型下载 ---
        with _status_lock:
            tasks_status[task_id]["message"] = "📥 Downloading LoRA / Preparing environment..."
            
        download_start = time.time()
        prepare_env(req.char_enable, req.char_lora, req.unet_choice, task_id=task_id)
        real_download_time = round(time.time() - download_start, 2)

        current_arch = UNET_CHOICES.get(req.unet_choice, {}).get("engine", "unknown")
        need_heavy_load = (_last_used_architecture != current_arch) or (_last_used_unet != req.unet_choice)

        if need_heavy_load:
            with _model_switch_lock:
                if _last_used_architecture != current_arch:
                    # 状态提示保留，卸载动作已移至 generate() 内部
                    with _status_lock:
                        tasks_status[task_id]["message"] = f"📦 Loading Base Model: {current_arch.upper()}..."
                    _last_used_architecture = current_arch
                    _last_used_unet = req.unet_choice
                elif _last_used_unet != req.unet_choice:
                    with _status_lock:
                        tasks_status[task_id]["message"] = "🔄 Swapping Model Checkpoint..."
                    _last_used_unet = req.unet_choice
        else:
            with _status_lock:
                tasks_status[task_id]["message"] = "🎨 GPU processing..."

        final_seed = req.seed if req.seed >= 0 else random.randint(0, 2**32 - 1)
        with _status_lock:
            tasks_status[task_id]["seed"] = final_seed
            tasks_status[task_id]["queue_start_time"] = time.time()
            tasks_status[task_id]["message"] = "⏳ Requesting GPU resources..."

        print(f"[GPU 探针] 🔴 发起 Request，申请 GPU 资源...", flush=True)

        ctx = contextvars.copy_context()

        def _run_generate():
            if request is not None:
                LocalContext.request.set(request)
            return generate(
                req.prompt, req.unet_choice, req.custom_w, req.custom_h,
                req.steps, req.cfg, final_seed, req.sampler, req.scheduler,
                req.char_enable, req.char_lora, req.char_strength,
                *req.base_lora_states,
                task_id=task_id,
            )

        # --- 阶段 2：核心执行（阻塞式，前端此时只看到进度条走动） ---
        image_path, real_workflow = ctx.run(_run_generate)

        # --- 阶段 3：死磕后处理（前端此时处于静默等待，避免出现“快好了”的错觉） ---
        from PIL import Image
        
        original_filename = os.path.basename(image_path)
        webp_filename = original_filename.replace(".png", ".webp")
        output_webp_path = os.path.join(OUTPUT, webp_filename)

        with Image.open(image_path) as img:
            metadata = img.info
            PROMPT_TAG = 0x0110
            EXTRA_METADATA_TAG = 0x010f
            output_metadata = img.getexif()
            
            for key, value in metadata.items():
                if key == 'prompt':
                    output_metadata[PROMPT_TAG] = "Prompt:" + value
                elif key == 'workflow':
                    output_metadata[EXTRA_METADATA_TAG] = "Workflow:" + value
                    EXTRA_METADATA_TAG -= 1

            if img.mode in ("RGBA", "P"):
                img = img.convert("RGB")

            # 🌟 性能妥协：method=4, quality=90，速度极快，肉眼无感
            img.save(
                output_webp_path, 
                "WEBP", 
                lossless=False, 
                quality=90, 
                method=4, 
                exif=output_metadata
            )

        try:
            os.remove(image_path)
        except Exception:
            pass

        image_url = f"/output/{webp_filename}"

        # --- 阶段 4：提取指标 ---
        final_queue = 0
        final_workflow = 0
        final_load_time = 0
        metric_file = f"/tmp/metrics_{task_id}.json"

        if os.path.exists(metric_file):
            try:
                with open(metric_file, "r") as f:
                    saved_metrics = json.load(f)
                    final_queue = saved_metrics.get("queue", 0)
                    final_workflow = saved_metrics.get("workflow", 0)
                    final_load_time = saved_metrics.get("load_time", 0)
                os.remove(metric_file)
            except Exception:
                pass

        # --- 阶段 5：高光时刻！直接抛出最终结果，抛弃一切中间态日志 ---
        with _status_lock:
            tasks_status[task_id] = {
                "status": "success",
                "image": image_url, 
                "filename": webp_filename,
                "seed": final_seed,
                "metrics": {
                    "download": real_download_time,
                    "queue": final_queue,
                    "workflow": final_workflow,
                    "load_time": final_load_time,
                    "steps": req.steps,
                    "resolved_seed": final_seed 
                },
                "anonymous": ip_token is None,
            }
            
    except Exception as e:
        with _status_lock:
            tasks_status[task_id] = {"status": "error", "message": str(e)}
    finally:
        _current_generation_task_id.set(None)

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

from fastapi import HTTPException

@app.post("/api/generate")
async def handle_generate(req: GenerateRequest, background_tasks: BackgroundTasks, request: Request):
    # 🌟 兜底拦截：单边最大 1368，防绕过
    if req.custom_w > 1536 or req.custom_h > 1536:
        raise HTTPException(status_code=400, detail="单边最大支持 1536 像素，请调整宽高。")
    
    if req.char_enable and req.char_lora:
        if not req.char_lora.endswith(".safetensors"):
            req.char_lora = f"{req.char_lora}.safetensors"

    task_id = str(uuid.uuid4())
    ip_token = request.headers.get("x-ip-token")
    tasks_status[task_id] = {"status": "queued", "message": "⏳ 任务已进入队列..."}

   # 👇 抛弃 BackgroundTasks，使用原生守护线程，直接与 HTTP 返回脱钩
    import threading
    t = threading.Thread(
        target=background_generate_task, 
        args=(task_id, req, ip_token, request),
        daemon=True
    )
    t.start()

    # 瞬间返回，绝不让前端等
    return {"task_id": task_id, "status": "queued"}

@app.get("/api/status/{task_id}")
async def check_status(task_id: str):
    status_data = tasks_status.get(task_id, {"status": "not_found"})
    if status_data.get("status") == "processing":
        metric_file = f"/tmp/metrics_{task_id}.json"
        if os.path.exists(metric_file):
            try:
                with open(metric_file, "r") as f:
                    file_data = json.load(f)
                    if "phase" in file_data:
                        status_data["message"] = file_data["phase"]
            except Exception:
                pass
    return status_data

@app.get("/api/models")
async def get_models():
    models_data = {}
    global_default_prompt = MODEL_CONFIG.get("default_prompt", "")

    for unet_name, unet_info in UNET_CHOICES.items():
        engine = unet_info["engine"]
        engine_name = MODEL_CONFIG["models"].get(engine, {}).get("name", engine)
        
        if " | " in unet_name:
            model_specific_name = unet_name.split(" | ", 1)[1]
        else:
            model_specific_name = unet_name
            
        display_unet_name = f"{engine_name} | {model_specific_name}"
        
        base_loras_for_engine = []
        for name, cfg in BASE_LORAS_CHOICES.items():
            if cfg["engine"] == engine:
                base_loras_for_engine.append({
                    "name": name,
                    "filename": cfg["filename"],
                    "defaultEnable": cfg.get("defaultEnable", False)
                })
        
        engine_cfg = MODEL_CONFIG["models"].get(engine, {})
        final_prompt = engine_cfg.get("default_prompt_override", global_default_prompt)

        # 🌟 核心新增：提取 TE 和 VAE 信息
        # TE 可能是列表，用 " + " 连接起来展示
        te_list = engine_cfg.get("local_te_filename", [])
        te_display = " + ".join([f.replace(".safetensors", "") for f in te_list]) if te_list else "Unknown"
        
        # VAE 通常是字符串
        vae_display = engine_cfg.get("local_vae_filename", "Unknown").replace(".safetensors", "")

        models_data[display_unet_name] = {
            "engine": engine,
            "defaults": engine_cfg.get("defaults", {}),
            "base_loras": base_loras_for_engine,
            "default_prompt": final_prompt,
            "te_name": te_display,   # 👈 传给前端
            "vae_name": vae_display  # 👈 传给前端
        }
        
    return {"models": models_data}

@app.get("/api/loras")
async def get_loras():
    load_char_loras_from_repo() 
    return {"loras": CHAR_LORA_CHOICES}

@app.get("/api/config")
async def get_sampler_config():
    return {
        "samplers": SAMPLERS,
        "schedulers": SCHEDULERS,
        "root_mapping": MODEL_CONFIG["json_keys"]["root_mapping"],
        "default_prompt": MODEL_CONFIG.get("default_prompt", ""),
    }

@app.get("/api/health")
async def health_check():
    return {"status": "ok"}

# ==========================================================================
# 📁 Static Files Sync & Mount
# ==========================================================================
from model_config import MODEL_CONFIG, SYNC_FROM_GUI_SPACE

UI_SPACE_URL = MODEL_CONFIG.get("ui_space_url", "").rstrip("/")
FRONTEND_CACHE_DIR = ROOT / "frontend_cache"  # 🌟 指向新文件夹

def sync_ui_assets():
    if not SYNC_FROM_GUI_SPACE:
        print("📦 [Frontend Cache] SYNC_FROM_GUI_SPACE = False，跳过同步，使用本地缓存。")
        return
        
    if not UI_SPACE_URL:
        print("⚠️ 未配置 ui_space_url，跳过前端同步。")
        return

    FRONTEND_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    files_to_sync = ["index.html", "style.css", "app.js"]
    
    print(f"🔄 正在从 {UI_SPACE_URL} 同步前端文件到 /frontend_cache...")
    for f in files_to_sync:
        try:
            res = requests.get(f"{UI_SPACE_URL}/{f.lstrip('/')}?t={int(time.time())}", timeout=10)
            if res.status_code == 200:
                (FRONTEND_CACHE_DIR / f).write_text(res.text, encoding="utf-8")
                print(f"✅ 已更新: {f}")
            else:
                print(f"⚠️ 同步 {f} 失败，状态码: {res.status_code}")
        except Exception as e:
            print(f"⚠️ 同步 {f} 网络异常: {e}")

# 1. 启动时执行同步
sync_ui_assets()

# ==========================================================
# 🌟 核心修改：必须在挂载根路径 "/" 之前，先挂载 /output 目录
# ==========================================================
app.mount("/output", StaticFiles(directory=str(OUTPUT)), name="output")

# 2. 挂载静态文件夹 (前端)
if (FRONTEND_CACHE_DIR / "index.html").exists():
    # html=True 会让根路径自动返回 index.html
    app.mount("/", StaticFiles(directory=str(FRONTEND_CACHE_DIR), html=True), name="frontend_cache")
else:
    @app.get("/")
    async def no_ui():
        return {"error": "UI files not found in /frontend_cache."}

# ==========================================================================
# 🚀 Server Startup
# ==========================================================================
if __name__ == "__main__":
    os.environ["GRADIO_SSR_MODE"] = "false"


    # 2. 启动 ZeroGPU 探针
    def trigger_gpu_scan():
        demo.launch(
            server_name="127.0.0.1",
            server_port=7861, 
            share=False,
            prevent_thread_lock=True,
        )
    
    print("🚀 Triggering ZeroGPU scan on 7861...")
    threading.Thread(target=trigger_gpu_scan, daemon=True).start()
    
    # 3. 稍微等一下扫描完成
    time.sleep(4) 

    # 4. 最后启动主 API 服务
    print("🌐 Starting main API Gateway on 7860...")
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")