"""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")