ImageGen / app.py
mrhangetsbetter's picture
Update app.py
c6383e8 verified
Raw
History Blame Contribute Delete
4.08 kB
import os
import sys
import site
# ================= 强力 CPU 环境全面拟真补丁 =================
if "--use-sage-attention" in sys.argv:
sys.argv.remove("--use-sage-attention")
if "--cpu" not in sys.argv:
sys.argv.append("--cpu")
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["FORCE_CPU"] = "1"
try:
import torch
# 1. 基础 CUDA 状态架空
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
torch.cuda.current_device = lambda: 0
# 2. 伪造 GPU 内存统计数据与 C++ 运行时属性
def mock_memory_stats(*args, **kwargs):
return {
'reserved_bytes.all.current': 0,
'reserved_bytes.all.peak': 0,
'allocated_bytes.all.current': 0,
'allocated_bytes.all.peak': 0
}
torch.cuda.memory_stats = mock_memory_stats
torch.cuda.memory_allocated = lambda *args, **kwargs: 0
torch.cuda.max_memory_allocated = lambda *args, **kwargs: 0
torch.cuda.memory_reserved = lambda *args, **kwargs: 0
torch.cuda.max_memory_reserved = lambda *args, **kwargs: 0
# 【核心修复】:拦截 mem_get_info,返回假的 (空闲显存, 总显存),直接给它 16GB 虚拟显存
torch.cuda.mem_get_info = lambda *args, **kwargs: (16 * 1024 * 1024 * 1024, 16 * 1024 * 1024 * 1024)
# 额外拦截设备属性查询,防止接下来报错
class MockDeviceProperties:
name = "Kaggle/HuggingFace CPU-Fake GPU"
major = 8
minor = 0
total_memory = 16 * 1024 * 1024 * 1024
torch.cuda.get_device_properties = lambda *args, **kwargs: MockDeviceProperties()
# 3. 构造设备代理类,保持 type 身份
_orig_device = torch.device
class TorchDeviceProxy:
def __new__(cls, *args, **kwargs):
if args and isinstance(args[0], str) and 'cuda' in args[0]:
return _orig_device('cpu')
try:
return _orig_device(*args, **kwargs)
except Exception:
return _orig_device('cpu')
@property
def __class__(self):
return _orig_device
TorchDeviceProxy.__name__ = _orig_device.__name__
TorchDeviceProxy.__module__ = _orig_device.__module__
TorchDeviceProxy.__qualname__ = _orig_device.__qualname__
torch.device = TorchDeviceProxy
print("🤖 [CPU Patch] Advanced TorchCuda/MemGetInfo mock deployed.")
except Exception as e:
print(f"⚠️ [CPU Patch] Error while generating proxy patch: {e}")
# ====================================================================
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)
def apply_sage_attention_patch():
print("--- [Runtime Patch] 🤖 CPU mode. SageAttention disabled. ---")
return "Skipped (CPU Mode)"
def dummy_gpu_decorator(func):
return func
@dummy_gpu_decorator
def dummy_gpu_for_startup():
print("--- [GPU Startup] Bypassed for CPU mode. ---")
return "Bypassed"
def main():
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}")
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()