Update download_api.py

download_api.py

@@ -35,6 +35,15 @@ app.add_middleware(
 # Global variable to track if processing is running
 processing_thread = None
 
+try:
+    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, check=True, shell=True)
+except subprocess.CalledProcessError as e:
+    print(f"Error installing flash-attn: {e}")
+    print("Continuing without flash-attn.")
+
+# Determine the device to use
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 def log_message(message):
     """Add a log message with timestamp"""
     timestamp = datetime.now().strftime("%H:%M:%S")