update app.py and model_worker.py

app.py

@@ -1,4 +1,3 @@
-import spaces
 import fire
 import subprocess
 import os
@@ -7,17 +6,6 @@ import signal
 import subprocess
 import atexit
 
-try:
-    import flash_attn
-except ImportError:
-
-    @spaces.GPU
-    def install_flash_attn():
-        os.system("pip install flash-attn==2.5.9.post1")
-
-    # install_flash_attn()
-    # import flash_attn
-
 
 def kill_processes_by_cmd_substring(cmd_substring):
     # execute `ps -ef` and obtain its output

model_worker.py

@@ -8,6 +8,7 @@
 A model worker executes the model.
 """
 import spaces
+import os
 import argparse
 import asyncio
 
@@ -310,6 +311,15 @@ class ModelWorker:
     @spaces.GPU
     @torch.inference_mode()
     def generate_stream(self, params):
+        try:
+            import flash_attn
+        except ImportError:
+
+            def install_flash_attn():
+                os.system("pip install flash-attn==2.5.9.post1")
+
+            install_flash_attn()
+            # import flash_attn
         system_message = params["prompt"][0]["content"]
         send_messages = params["prompt"][1:]
         max_input_tiles = params["max_input_tiles"]