File size: 3,829 Bytes
1875ee2
c2bcd10
 
1875ee2
c2bcd10
d9220fa
 
 
 
1875ee2
 
 
 
c2bcd10
1875ee2
c2bcd10
1875ee2
 
 
 
c2bcd10
 
1875ee2
 
 
 
 
 
 
 
 
 
 
 
 
 
c2bcd10
1875ee2
 
 
 
c2bcd10
1875ee2
 
0bce0ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c2bcd10
d52ffc0
1875ee2
d52ffc0
61e9353
c2bcd10
d52ffc0
 
 
d9220fa
 
 
 
1875ee2
c2bcd10
1875ee2
 
 
 
 
d9220fa
d52ffc0
 
 
 
1875ee2
d52ffc0
1875ee2
 
 
 
c2bcd10
 
1875ee2
 
 
c2bcd10
1875ee2
c2bcd10
1875ee2
 
 
 
c2bcd10
 
 
1875ee2
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
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 apply_sage_attention_patch():
    global SAGE_PATCH_APPLIED
    if SAGE_PATCH_APPLIED:
        return "SageAttention patch already applied."

    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():
    try:
        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."
    except BaseException as e:
        err_msg = str(e)
        if "uncorrectable ECC error" in err_msg or "cudaErrorECCUncorrectable" in err_msg:
            print("\n" + "="*80)
            print(f"🚨 [Fatal GPU Error] Captured uncorrectable ECC error during inference: {err_msg}")
            print("🚨 Terminating process to trigger an automatic container restart...")
            print("="*80 + "\n")
            os._exit(1)
        raise e


def main():
    from comfy_integration import setup as setup_comfyui
    from 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} ---")

    print("--- Initiating GPU Startup Check & SageAttention Patch Verification ---")
    try:
        dummy_gpu_for_startup()
    except Exception as e:
        print(f"--- [GPU Startup] ⚠️ Warning: Startup check failed: {e} ---")

    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

    print(f"✅ Working directory is stable: {os.getcwd()}")

    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()