Update app.py

app.py

@@ -154,27 +154,42 @@ def install_flash_attn():
             
         logging.info(f"Detected CUDA version: {cuda_version}")
         
-        # CUDA 버전별 wheel 파일 선택
-        if cuda_version.startswith("12.1"):
-            flash_attn_url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.5/flash_attn-2.7.5+cu121torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
-        elif cuda_version.startswith("11.8"):
-            flash_attn_url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu11torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
-        else:
-            logging.warning(f"Unsupported CUDA version: {cuda_version}, skipping flash-attn installation")
-            return False
+        try:
+            import flash_attn
+            logging.info("flash-attn already installed")
+            return True
+        except ImportError:
+            logging.info("Installing flash-attn...")
             
-        subprocess.run(
-            ["pip", "install", flash_attn_url],
-            check=True,
-            capture_output=True
-        )
+        # CUDA 12.1용 직접 설치 시도
+        try:
+            subprocess.run(
+                ["pip", "install", "flash-attn", "--no-build-isolation"],
+                check=True,
+                capture_output=True
+            )
+            logging.info("flash-attn installed successfully!")
+            return True
+        except subprocess.CalledProcessError:
+            logging.warning("Failed to install flash-attn via pip, skipping...")
+            return False
             
-        logging.info("flash-attn installed successfully!")
-        return True
     except Exception as e:
         logging.warning(f"Failed to install flash-attn: {e}")
         return False
 
+# ... (나머지 코드는 동일) ...
+
+# 서버 설정으로 실행 부분만 수정
+demo.queue(max_size=20).launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=True,
+    show_api=True,
+    show_error=True,
+    max_threads=2  # concurrency_count 대신 max_threads 사용
+)
+
 def initialize_system():
     optimize_gpu_settings()
     has_flash_attn = install_flash_attn()