Spaces:
Running on Zero
Running on Zero
File size: 4,520 Bytes
8a28a8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import spaces
import importlib.util
import os
import sys
import site
from core.runtime_config import CONFIG
sage_mode = os.getenv("IMAGEGEN_USE_SAGE_ATTENTION", "auto").strip().lower()
sage_available = importlib.util.find_spec("sageattention") is not None
use_sage_attention = sage_mode in {"1", "true", "yes", "on"} or (
sage_mode == "auto" and sage_available
)
if use_sage_attention and "--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.")
# Keep ComfyUI code isolated while pointing its model/input/output directories
# at this Space. These arguments are consumed when ComfyUI is imported.
if "--base-directory" not in sys.argv:
sys.argv.extend(["--base-directory", APP_DIR])
SAGE_PATCH_APPLIED = False
def apply_sage_attention_patch():
global SAGE_PATCH_APPLIED
if SAGE_PATCH_APPLIED:
return "SageAttention patch already applied."
if not use_sage_attention:
return "SageAttention disabled or unavailable; using the default attention backend."
try:
from comfy import model_management
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:
SAGE_PATCH_APPLIED = False
msg = "--- [Runtime Patch] ⚠️ sageattention package not found. 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():
os.chdir(APP_DIR)
from comfy_integration import setup as setup_comfyui
from imagegen_utils.app_utils import load_ipadapter_presets
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} ---")
if CONFIG.enable_startup_gpu_probe:
print("--- Initiating optional GPU startup check ---")
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("--- Loading IPAdapter presets ---")
load_ipadapter_presets()
print("--- ✅ IPAdapter setup complete. ---")
print("--- Environment configured. Proceeding with module imports. ---")
from ui.layout import build_ui
from ui.events import attach_event_handlers
if CONFIG.enable_mcp:
import mcp_tools as mcp
print(f"✅ Loaded MCP module with tools: {[fn.__name__ for fn in mcp.MCP_FUNCTIONS]}")
else:
print("ℹ️ MCP is disabled by IMAGEGEN_ENABLE_MCP.")
print(f"✅ Working directory is stable: {os.getcwd()}")
demo = build_ui(attach_event_handlers)
print(
"--- Launching Gradio Interface "
f"(GPU concurrency={CONFIG.gpu_concurrency}, queue={CONFIG.queue_max_size}, "
f"MCP={CONFIG.enable_mcp}) ---"
)
demo.queue(
default_concurrency_limit=1,
max_size=CONFIG.queue_max_size,
status_update_rate="auto",
).launch(mcp_server=CONFIG.enable_mcp)
if __name__ == "__main__":
main()
|