Update src/florence/model.py

src/florence/model.py

@@ -1,11 +1,11 @@
 # Importing necessary libraries
-import spaces
 import sys
 import subprocess
 from typing import Optional
 from PIL import Image
 import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
+import spaces
 import gradio as gr
 
 # Local imports
@@ -13,12 +13,40 @@ from src.logger import logging
 from src.exception import CustomExceptionHandling
 
 
-# Install the required dependencies
-subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
-)
+# Install the required dependencies (flash-attn)
+# Check if flash_attn is already installed before attempting installation
+try:
+    import flash_attn
+    logging.info("flash_attn is already installed.")
+except ImportError:
+    logging.info("flash_attn not found. Attempting to install...")
+    try:
+        # Set environment variable to skip CUDA build if needed
+        import os
+        env = os.environ.copy()
+        env["FLASH_ATTENTION_SKIP_CUDA_BUILD"] = "TRUE"
+        
+        result = subprocess.run(
+            ["pip", "install", "flash-attn", "--no-build-isolation"],
+            env=env,
+            capture_output=True,
+            text=True,
+            check=False
+        )
+        
+        if result.returncode == 0:
+            logging.info("flash_attn installed successfully.")
+        else:
+            logging.warning(f"flash_attn installation may have failed: {result.stderr}")
+            # Try importing again to verify
+            try:
+                import flash_attn
+                logging.info("flash_attn import successful after installation.")
+            except ImportError:
+                logging.error("flash_attn installation failed. The model may not work correctly.")
+    except Exception as e:
+        logging.warning(f"Error during flash_attn installation: {e}")
+        logging.warning("Continuing without flash_attn. The model may require it.")
 
 # Load model and processor from Hugging Face
 model_id = "microsoft/Florence-2-large-ft"
@@ -38,9 +66,7 @@ except Exception as e:
 
 
 @spaces.GPU
-def run_example(
-    task_prompt: str, image: Image.Image, text_input: Optional[str] = None
-) -> str:
+def run_example(task_prompt: str, image: Image.Image, text_input: Optional[str] = None) -> str:
     """
     Runs an example using the given task prompt and image.