Update src/florence/model.py

src/florence/model.py

@@ -13,40 +13,12 @@ from src.logger import logging
 from src.exception import CustomExceptionHandling
 
 
-# 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.")
+# Install the required dependencies
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 
 # Load model and processor from Hugging Face
 model_id = "microsoft/Florence-2-large-ft"
@@ -66,7 +38,9 @@ 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.