ImageGen / app.py
mrhangetsbetter's picture
Update app.py
4bbc5ee verified
Raw
History Blame
4.44 kB
import os
import sys
import site
# ================= 强力 CPU 猴子补丁 (Monkey Patch) =================
# 必须在导入任何 comfy 或 torch 模块之前执行!
# 1. 强行清理和重写启动参数
if "--use-sage-attention" in sys.argv:
sys.argv.remove("--use-sage-attention")
if "--cpu" not in sys.argv:
sys.argv.append("--cpu")
# 2. 伪造 CUDA 环境变量
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["FORCE_CPU"] = "1"
# 3. 拦截并重写 torch.cuda 的核心方法,彻底废除其硬件检测能力
try:
import torch
# 强制让 torch 认为 cuda 不可用
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
# 劫持 current_device 和 get_device_properties,让其不触发底层 C++ 的 _cuda_init()
def mock_current_device():
raise AssertionError("CUDA is not available in CPU-only mode.")
torch.cuda.current_device = mock_current_device
# 关键:直接重写 ComfyUI 报错的那一行代码所依赖的底层逻辑
# 模拟一个假的设备对象或直接让 get_torch_device 返回 cpu
print("🤖 [CPU Patch] Successfully mocked torch.cuda to prevent driver crash.")
except Exception as e:
print(f"⚠️ [CPU Patch] Failed to pre-patch torch: {e}")
# ====================================================================
# 尝试导入 spaces(ZeroGPU 专用,CPU 环境下保留即可)
try:
import spaces
except ImportError:
spaces = None
APP_DIR = os.path.dirname(os.path.abspath(__file__))
if APP_DIR not in sys.path:
sys.path.insert(0, APP_DIR)
# 动态劫持即将被导入的 comfy.model_management 的硬件获取函数
# 这是一个更激进的拦截:直接干涉 Python 的模块系统
class ComfyManagementHacker:
def __init__(self):
self.patched = False
def find_spec(self, fullname, path, target=None):
# 当系统尝试导入 comfy.model_management 时,我们提前介入
if fullname == "comfy.model_management" and not self.patched:
self.patched = True
print("🎯 [Interceptor] Caught comfy.model_management loading! Injecting CPU device detour...")
# 提前把 torch.cuda 再次阉割
import torch
torch.device = self.make_cpu_device_wrapper(torch.device)
return None
def make_cpu_device_wrapper(self, original_device):
def wrapped_device(*args, **kwargs):
# 任何时候如果尝试创建 cuda 设备,强行扭转为 cpu
if args and ('cuda' in str(args[0])):
return original_device('cpu')
return original_device(*args, **kwargs)
return wrapped_device
sys.meta_path.insert(0, ComfyManagementHacker())
def apply_sage_attention_patch():
print("--- [Runtime Patch] 🤖 CPU mode. SageAttention disabled. ---")
return "Skipped (CPU Mode)"
def dummy_gpu_decorator(func):
return func # 完全废除 @spaces.GPU 带来的潜在干扰
@dummy_gpu_decorator
def dummy_gpu_for_startup():
print("--- [GPU Startup] Bypassed for CPU mode. ---")
return "Bypassed"
def main():
# 在这里,当 setup_comfyui 内部导入 comfy.model_management 时,
# 我们的拦截器或上面的 torch 补丁会确保它不会去触发 C++ 的 _cuda_init()
from comfy_integration import setup as setup_comfyui
from utils.app_utils import load_ipadapter_presets
print("--- [Setup] Starting ComfyUI initialization ---")
try:
setup_comfyui.initialize_comfyui()
except Exception as e:
print(f"🚨 [Init Error] ComfyUI initialization failed: {e}")
print("💡 If it still mentions 'RuntimeError: Found no NVIDIA driver', ComfyUI's specific script bypassed the patch. You may need a GPU instance.")
raise e
print("--- [Setup] Reloading site-packages... ---")
try:
site.main()
except Exception:
pass
print("--- Starting Application Setup ---")
load_ipadapter_presets()
from ui.layout import build_ui
from ui.events import attach_event_handlers
demo = build_ui(attach_event_handlers)
print("--- Launching Gradio Interface ---")
demo.queue().launch(server_name="0.0.0.0", server_port=7860)
if __name__ == "__main__":
main()