File size: 4,877 Bytes
5b29993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
129
import spaces
import os
import sys
import requests
import site

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():
    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 utils.app_utils import print_welcome_message
    from scripts import build_sage_attention

    print_welcome_message()
    
    print("--- [Setup] Attempting to build and install SageAttention... ---")
    try:
        build_sage_attention.install_sage_attention()
        print("--- [Setup] βœ… SageAttention installation process finished. ---")
    except Exception as e:
        print(f"--- [Setup] ❌ SageAttention installation failed: {e}. Continuing with default attention. ---")


    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} ---")

    from comfy_integration import setup as setup_comfyui
    from utils.app_utils import (
        build_preprocessor_model_map,
        build_preprocessor_parameter_map,
        load_ipadapter_presets
    )
    from core import shared_state
    from core.settings import ALL_MODEL_MAP, ALL_FILE_DOWNLOAD_MAP

    def check_all_model_urls_on_startup():
        print("--- [Setup] Checking all model URL validity (one-time check) ---")
        for display_name, model_info in ALL_MODEL_MAP.items():
            repo_id, filename, _, _ = model_info
            if not repo_id: continue

            download_info = ALL_FILE_DOWNLOAD_MAP.get(filename, {})
            repo_file_path = download_info.get('repository_file_path', filename)
            url = f"https://huggingface.co/{repo_id}/resolve/main/{repo_file_path}"
            
            try:
                response = requests.head(url, timeout=5, allow_redirects=True)
                if response.status_code >= 400:
                    print(f"❌ Invalid URL for '{display_name}': {url} (Status: {response.status_code})")
                    shared_state.INVALID_MODEL_URLS[display_name] = True
            except requests.RequestException as e:
                print(f"❌ URL check failed for '{display_name}': {e}")
                shared_state.INVALID_MODEL_URLS[display_name] = True
        print("--- [Setup] βœ… Finished checking model URLs. ---")

    print("--- Starting Application Setup ---")
    
    setup_comfyui.initialize_comfyui()

    check_all_model_urls_on_startup()
    
    print("--- Building ControlNet preprocessor maps ---")
    from core.generation_logic import build_reverse_map
    build_reverse_map()
    build_preprocessor_model_map()
    build_preprocessor_parameter_map()
    print("--- βœ… ControlNet preprocessor setup complete. ---")

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