Spaces:
Running on Zero
Running on Zero
"""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")
"""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:
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():