Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,073 Bytes
1875ee2 9de6815 1875ee2 0f16a33 1875ee2 0f16a33 9de6815 1875ee2 9de6815 1875ee2 9de6815 1875ee2 9de6815 1875ee2 9de6815 1875ee2 9de6815 1875ee2 9de6815 1875ee2 9de6815 1875ee2 9de6815 1875ee2 9de6815 |
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 |
import os
import subprocess
import sys
def run_command(command, env=None):
"""
Runs a command with a specified environment, prints its output,
and raises an exception on failure.
"""
print(f"🚀 Running command: {' '.join(command)}")
result = subprocess.run(
command,
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding='utf-8',
errors='replace'
)
if result.stdout:
print("--- Pip Output ---")
print(result.stdout.strip())
print("------------------")
if result.returncode != 0:
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 ---")
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']}'")
install_command = [
sys.executable,
"-m",
"pip",
"install",
"sageattention==2.2.0",
"--no-build-isolation",
]
try:
run_command(install_command, env=build_env)
print("🎉 SageAttention installed successfully via pip!")
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
except Exception as e:
print(f"❌ An unexpected error occurred during pip installation: {e}")
raise e
if __name__ == "__main__":
try:
install_sage_attention()
except Exception:
print("\nInstallation script finished with an error.")
sys.exit(1) |