Upload folder using huggingface_hub

comfy_integration/setup.py

@@ -20,7 +20,7 @@ def initialize_comfyui():
 
     print("--- Cloning ComfyUI Repository ---")
     if not os.path.exists(COMFYUI_TEMP_DIR):
-        os.system(f"git clone https://github.com/comfyanonymous/ComfyUI {COMFYUI_TEMP_DIR}")
+        os.system(f"git clone https://github.com/comfy-Org/ComfyUI {COMFYUI_TEMP_DIR}")
         print("✅ ComfyUI repository cloned.")
     else:
         print("✅ ComfyUI repository already exists.")

scripts/build_sage_attention.py

@@ -1,66 +1,99 @@
 import os
 import subprocess
 import sys
+import textwrap
 
-def run_command(command, env=None):
-    """
-    Runs a command with a specified environment, prints its output, 
-    and raises an exception on failure.
-    """
+REPO_URL = "https://github.com/thu-ml/SageAttention.git"
+REPO_DIR = "SageAttention"
+
+def run_command(command, cwd=None, env=None):
     print(f"🚀 Running command: {' '.join(command)}")
     result = subprocess.run(
         command,
+        cwd=cwd,
         env=env,
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
-        text=True,
-        encoding='utf-8',
-        errors='replace'
+        text=True
     )
-
-    if result.stdout:
-        print("--- Pip Output ---")
-        print(result.stdout.strip())
-        print("------------------")
-
+    
     if result.returncode != 0:
+        print(result.stdout)
         raise subprocess.CalledProcessError(result.returncode, command)
 
-def install_sage_attention():
-    """
-    Installs the sageattention package from PyPI using pip, ensuring the
-    correct CUDA architecture is set for any potential on-the-fly compilation.
-    """
-    print("--- [SageAttention Install] Starting installation using pip ---")
+def patch_setup_py(setup_py_path):
+    print(f"--- [SageAttention Build] Applying patches to {setup_py_path} ---")
 
-    build_env = os.environ.copy()
-    build_env["TORCH_CUDA_ARCH_LIST"] = "9.0"
-    print(f"🔧 Setting build environment variable: TORCH_CUDA_ARCH_LIST='{build_env['TORCH_CUDA_ARCH_LIST']}'")
+    with open(setup_py_path, 'r', encoding='utf-8') as f:
+        content = f.read()
 
-    install_command = [
-        sys.executable,
-        "-m",
-        "pip",
-        "install",
-        "sageattention==2.2.0",
-        "--no-build-isolation",
-    ]
+    original_cxx_flags = 'CXX_FLAGS = ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"]'
+    modified_cxx_flags = 'CXX_FLAGS = ["-g", "-O3", "-std=c++17", "-DENABLE_BF16"]'
+    
+    if original_cxx_flags in content:
+        content = content.replace(original_cxx_flags, modified_cxx_flags)
+        print("🔧 Patch 1/1: Removed '-fopenmp' and '-lgomp' from CXX_FLAGS.")
+    else:
+        print("⚠️ Patch 1/1: CXX_FLAGS line not found as expected. It might have been changed upstream. Skipping.")
+
+    with open(setup_py_path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    
+    print("✅ Patches applied successfully.")
+
+
+def install_sage_attention():
+    print("--- [SageAttention Build] Checking environment ---")
+    
+    if os.path.isdir(REPO_DIR):
+        print(f"✅ Directory '{REPO_DIR}' already exists, assuming SageAttention is installed. Skipping build.")
+        return
 
+    print(f"⏳ Directory '{REPO_DIR}' not found. Starting a fresh installation of SageAttention.")
+    
     try:
-        run_command(install_command, env=build_env)
-        print("🎉 SageAttention installed successfully via pip!")
+        print(f"--- [SageAttention Build] Step 1/3: Cloning repository ---")
+        run_command(["git", "clone", REPO_URL])
+        print("✅ Repository cloned successfully.")
+
+        print(f"--- [SageAttention Build] Step 2/3: Patching setup.py ---")
+        setup_py_path = os.path.join(REPO_DIR, "setup.py")
+        patch_setup_py(setup_py_path)
+
+        print(f"--- [SageAttention Build] Step 3/3: Compiling and installing ---")
+        
+        build_env = os.environ.copy()
+        build_env.update({
+            "TORCH_CUDA_ARCH_LIST": "9.0",
+            "EXT_PARALLEL": "4",
+            "NVCC_APPEND_FLAGS": "--threads 8",
+            "MAX_JOBS": "32"
+        })
+        print("🔧 Setting build environment variables:")
+        print(f"   - TORCH_CUDA_ARCH_LIST='{build_env['TORCH_CUDA_ARCH_LIST']}'")
+        print(f"   - EXT_PARALLEL={build_env['EXT_PARALLEL']}")
+        print(f"   - NVCC_APPEND_FLAGS='{build_env['NVCC_APPEND_FLAGS']}'")
+        print(f"   - MAX_JOBS={build_env['MAX_JOBS']}")
+
+        install_command = [sys.executable, "setup.py", "install"]
+        
+        run_command(install_command, cwd=REPO_DIR, env=build_env)
+
+        print("🎉 SageAttention compiled and installed successfully! ---")
+
+    except FileNotFoundError:
+        print("❌ ERROR: 'git' command not found. Please ensure Git is installed in your environment.")
+        sys.exit(1)
     except subprocess.CalledProcessError as e:
-        print(f"❌ SageAttention installation via pip failed (Exit Code: {e.returncode}).")
-        print("   The application will continue with the default attention mechanism.")
-        raise e
+        print(f"❌ Command failed with return code: {e.returncode}")
+        print(f"❌ Command: {' '.join(e.cmd)}")
+        print("❌ SageAttention installation failed. Please check the logs above for details.")
+        sys.exit(1)
     except Exception as e:
-        print(f"❌ An unexpected error occurred during pip installation: {e}")
-        raise e
-
+        print(f"❌ An unknown error occurred: {e}")
+        sys.exit(1)
 
 if __name__ == "__main__":
-    try:
-        install_sage_attention()
-    except Exception:
-        print("\nInstallation script finished with an error.")
-        sys.exit(1)
+    if os.path.isdir(REPO_DIR):
+         print(f"Note: To force a rebuild, please delete the '{REPO_DIR}' directory first.")
+    install_sage_attention()