import spaces import os import sys import site if "--use-sage-attention" not in sys.argv: sys.argv.append("--use-sage-attention") print("🚀 [SageAttention] Injected '--use-sage-attention' into sys.argv.") APP_DIR = os.path.dirname(os.path.abspath(__file__)) if APP_DIR not in sys.path: sys.path.insert(0, APP_DIR) print(f"✅ Added project root '{APP_DIR}' to sys.path.") SAGE_PATCH_APPLIED = False def _preload_cuda_runtime(): """Preload CUDA runtime libs so sageattention's DT_NEEDED (libcudart.so.13) resolves. torch 2.13 静态链接了 cudart（容器里没有独立的 libcudart.so.13 文件）， 而预编译的 sageattention wheel 是动态链接 libcudart.so.13 的。 这里用 ctypes.RTLD_GLOBAL 把 CUDA runtime 提前载入进程， 之后 import sageattention 时 DT_NEEDED 按 SONAME 命中已加载的库，不再找文件路径。 """ import ctypes import glob as _glob candidates = [] try: import torch as _torch _tdir = os.path.dirname(_torch.__file__) candidates += _glob.glob(os.path.join(_tdir, "lib", "libcudart.so*")) candidates += _glob.glob(os.path.join(_tdir, "..", "nvidia", "cuda_runtime", "lib", "libcudart.so*")) except Exception: pass candidates += _glob.glob(os.path.join(APP_DIR, "core", "cuda_libs", "libcudart.so*")) loaded = [] for p in candidates: try: ctypes.CDLL(p, mode=ctypes.RTLD_GLOBAL) loaded.append(os.path.basename(p)) except OSError: continue if loaded: print(f"--- [Runtime Patch] Preloaded CUDA runtime libs: {loaded} ---") return bool(loaded) def apply_sage_attention_patch(): global SAGE_PATCH_APPLIED if SAGE_PATCH_APPLIED: return "SageAttention patch already applied." try: from comfy import model_management _preload_cuda_runtime() import sageattention print("--- [Runtime Patch] sageattention package found. Applying patch... ---") model_management.sage_attention_enabled = lambda: True model_management.pytorch_attention_enabled = lambda: False SAGE_PATCH_APPLIED = True return "✅ Successfully enabled SageAttention." except ImportError as e: SAGE_PATCH_APPLIED = False msg = f"--- [Runtime Patch] ⚠️ sageattention import failed ({type(e).__name__}: {e}). Continuing with default attention. ---" print(msg) return msg except Exception as e: SAGE_PATCH_APPLIED = False msg = f"--- [Runtime Patch] ❌ An error occurred while applying SageAttention patch: {e} ---" print(msg) return msg @spaces.GPU def dummy_gpu_for_startup(): print("--- [GPU Startup] Dummy function for startup check initiated. ---") patch_result = apply_sage_attention_patch() print(f"--- [GPU Startup] {patch_result} ---") print("--- [GPU Startup] Startup check passed. ---") return "Startup check passed." def main(): from comfy_integration import setup as setup_comfyui print("--- [Setup] Starting ComfyUI initialization ---") setup_comfyui.initialize_comfyui() print("--- [Setup] Applying SageAttention Runtime Patch ---") patch_result = apply_sage_attention_patch() print(f"--- [Setup] {patch_result} ---") print("--- [Setup] Reloading site-packages to detect newly installed packages... ---") try: site.main() print("--- [Setup] ✅ Site-packages reloaded. ---") except Exception as e: print(f"--- [Setup] ⚠️ Warning: Could not fully reload site-packages: {e} ---") print("--- Initiating GPU Startup Check & SageAttention Patch Verification ---") try: dummy_gpu_for_startup() except BaseException as e: err_msg = f"{type(e).__name__}: {str(e)}" print(f"--- [GPU Startup] ⚠️ Warning: Startup check failed: {err_msg} ---") print("--- Starting Application Setup ---") print("--- Environment configured. Proceeding with module imports. ---") from ui.layout import build_ui from ui.events import attach_event_handlers import mcp_tools as mcp print(f"✅ Loaded MCP module with tools: {[fn.__name__ for fn in mcp.MCP_FUNCTIONS]}") print(f"✅ Working directory is stable: {os.getcwd()}") demo = build_ui(attach_event_handlers) print("--- Launching Gradio Interface ---") # ===== 访问门禁 v4: 只拦网页首页, API 放行 ===== try: import base64 as _b64 import gradio.blocks as _gb from starlette.middleware.base import BaseHTTPMiddleware as _BHM from starlette.responses import PlainTextResponse as _PTR _GATE_USER, _GATE_PASS = "h3", "v0ChK59bPTEIZTJTSbA6" class _UIGate(_BHM): async def dispatch(self, request, call_next): p = (request.url.path or "").rstrip("/") if p in ("", "/"): # 只挡网页 UI 首页 hdr = request.headers.get("authorization", "") or "" ok = False if hdr.lower().startswith("basic "): try: _u, _pw = _b64.b64decode(hdr.split(" ", 1)[1]).decode("utf-8").split(":", 1) ok = (_u == _GATE_USER and _pw == _GATE_PASS) except Exception: ok = False if not ok: return _PTR("Locked: password required.", status_code=401, headers={"WWW-Authenticate": 'Basic realm="H3"'}) return await call_next(request) # 其余(API/静态资源)一律放行 # 关键: launch() 内部会 App.create_app() 重建 app, 必须 patch 工厂才能挂到最终 app 上 _App = getattr(_gb, "App", None) if _App is not None and not getattr(_App, "_h3gate", False): _orig = _App.create_app def _create_app_with_gate(*a, **kw): _app = _orig(*a, **kw) try: _app.add_middleware(_UIGate) print("[gate] OK: 门禁已挂到最终 app") except Exception as _e2: print("[gate] 挂载失败: %s" % _e2) return _app _App.create_app = staticmethod(_create_app_with_gate) _App._h3gate = True print("[gate] v4 已注入 (App=%s)" % _App) except Exception as _ge: print("[gate] 门禁未启用(不影响空间): %s" % _ge) # ===== 访问门禁结束 ===== demo.queue().launch(mcp_server=True) if __name__ == "__main__": main()