app.py

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