Update app.py

app.py

@@ -4,20 +4,21 @@
 # in the Software without restriction, including without limitation the rights
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 # copies of the Software, and to permit persons to whom the Software is
-#import spaces
-#import subprocess
+import spaces
+import subprocess
 import os
-#subprocess.run(['sh', './torch.sh'])
 
-os.putenv('PYTORCH_NVML_BASED_CUDA_CHECK','1')
-os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
-alloc_conf_parts = [
-    'expandable_segments:True',
-    'pinned_use_background_threads:True'  # Specific to pinned memory.
-]
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(alloc_conf_parts)
+@spaces.GPU(required=True)
+def install_flashattn():
+    subprocess.run(['sh', './flashattn.sh'])
+    
+install_flashattn()
+
+os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
+os.environ['TORCH_LINALG_PREFER_CUSOLVER'] = '1'
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,pinned_use_background_threads:True'
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
-os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
+os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1'
 
 import gradio as gr
 import numpy as np