Upload 8 files

app.py

@@ -366,6 +366,33 @@ class AppState:
 
         print("✓ RF-DETR model loaded")
 
+    def preload_all_models(self):
+        """Preload both detection and LLM models into VRAM at startup."""
+        print("=" * 60)
+        print("Preloading all models into VRAM...")
+        print("=" * 60)
+        
+        # Load detection model
+        print("\n[1/2] Loading RF-DETR detection model...")
+        self.load_model()
+        
+        # Load LLM model
+        if self.config.get('use_llm'):
+            print("\n[2/2] Loading MedGemma LLM model...")
+            try:
+                model_size = "4B"  # Default to 4B model
+                generator = self.get_text_generator(model_size)
+                hf_token = self.config.get('hf_token')
+                generator.load_model(hf_token)
+                print("✓ MedGemma model loaded and ready")
+            except Exception as e:
+                print(f"⚠️ Warning: Could not preload LLM model: {e}")
+                print("LLM will be loaded on first use instead")
+        
+        print("\n" + "=" * 60)
+        print("✓ All models loaded and ready in VRAM!")
+        print("=" * 60 + "\n")
+
     def get_text_generator(self, model_size: str = "4B") -> TextGenerator:
         """Get or create text generator."""
         # Determine model ID based on size selection
@@ -403,8 +430,9 @@ def create_detection_interface():
             return None, "Please upload an image."
 
         try:
-            # Load model if needed
-            app_state.load_model()
+            # Models are preloaded at startup, but check just in case
+            if app_state.model is None:
+                app_state.load_model()
 
             # Run detection
             detections = app_state.model.predict(image, threshold=threshold)
@@ -456,6 +484,7 @@ def create_detection_interface():
                     try:
                         generator = app_state.get_text_generator(model_size)
                         hf_token = app_state.config.get('hf_token')
+                        # Model is already preloaded, just generate
                         llm_description = generator.generate(description, image=annotated, hf_token=hf_token)
                         description = llm_description
                     except Exception as e:
@@ -560,6 +589,13 @@ def main():
     # Ensure results directory exists
     os.makedirs(app_state.config.get('results_dir'), exist_ok=True)
 
+    # Preload all models into VRAM
+    try:
+        app_state.preload_all_models()
+    except Exception as e:
+        print(f"⚠️ Warning: Failed to preload models: {e}")
+        print("Models will be loaded on first use instead")
+
     # Create and launch the interface
     demo = create_detection_interface()