Hugging Face's logo

Spaces:
RioShiina
/
ImageGen-NewBie-Image
Running on Zero

App Files Community
ImageGen-NewBie-Image
File size: 4,713 Bytes 
41978ed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
 
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():
    """Applies the SageAttention patch for optimized performance."""
    global SAGE_PATCH_APPLIED
    if SAGE_PATCH_APPLIED:
        msg = "--- [Runtime Patch] ✅ SageAttention patch is already applied. ---"
        print(msg)
        return msg

    try:
        import comfy.model_management as model_management
        from sage_attention.patch import apply_sage_attention
        apply_sage_attention(model_management)
        SAGE_PATCH_APPLIED = True
        msg = "--- [Runtime Patch] ✅ SageAttention patch successfully applied. ---"
        print(msg)
        return msg
    except ImportError:
        SAGE_PATCH_APPLIED = False
        msg = "--- [Runtime Patch] ⚠️ SageAttention not found or failed to install. Continuing with default attention. ---"
        print(msg)
        return msg
    except Exception as e:
        SAGE_PATCH_APPLIED = False
        msg = f"--- [Runtime Patch] ❌ SageAttention patch failed: {e}. Continuing with default attention. ---"
        print(msg)
        return msg

@spaces.GPU
def dummy_gpu_for_startup():
    """This function is decorated with @spaces.GPU to ensure a GPU is allocated on startup."""
    print("--- [GPU Startup] Dummy function for startup check passed. GPU should be allocated. ---")
    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 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():
            _, components, _, _ = model_info
            if not components: continue
            
            for filename in components.values():
                download_info = ALL_FILE_DOWNLOAD_MAP.get(filename, {})
                repo_id = download_info.get('repo_id')
                if not repo_id: continue
                
                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}' component '{filename}': {url} (Status: {response.status_code})")
                        shared_state.INVALID_MODEL_URLS[display_name] = True
                        break 
                except requests.RequestException as e:
                    print(f"❌ URL check failed for '{display_name}' component '{filename}': {e}")
                    shared_state.INVALID_MODEL_URLS[display_name] = True
                    break
        print("--- [Setup] ✅ Finished checking model URLs. ---")

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

    apply_sage_attention_patch()

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