Update app.py

app.py

@@ -25,24 +25,31 @@ MODEL_REGISTRY = {
         "grounding_dino": "grounding_dino_detect_model"
     }
 }
-LOADED_MODELS = {}
 
+# --- CORRECTED MODEL MANAGEMENT ---
+# Caching is removed to prevent errors from stateful models.
+# This function now loads a fresh model for each analysis request.
 def get_model(task, architecture):
-    """Lazily loads a model based on user selection and caches it."""
+    """
+    Loads a fresh model instance based on user selection.
+    This prevents stateful changes from one run affecting the next.
+    """
     try:
+        # For Zero-Shot, the architecture is always 'grounding_dino'
+        if task == "Zero-Shot Detection":
+            architecture = "grounding_dino"
+            
         model_name = MODEL_REGISTRY[task][architecture]
-        if model_name not in LOADED_MODELS:
-            print(f"Loading model for the first time: {model_name}")
-            LOADED_MODELS[model_name] = ibbi.create_model(model_name, pretrained=True)
-            print("Model loaded successfully!")
-        return LOADED_MODELS[model_name]
+        print(f"Loading a fresh model instance: {model_name}")
+        model = ibbi.create_model(model_name, pretrained=True)
+        print("Model loaded successfully!")
+        return model
     except KeyError as e:
         raise gr.Error(f"Model lookup failed. Task: '{task}', Arch: '{architecture}'. Error: {e}")
     except Exception as e:
         raise gr.Error(f"Failed to load model. Please check the model name and your connection. Error: {e}")
 
 # --- Visualization and Drawing Functions ---
-# Note: The global font object has been removed from here.
 
 def draw_yolo_predictions(image, results, font, color="red"):
     """Draws YOLO predictions on an image with a dynamically sized font."""
@@ -100,23 +107,20 @@ def visualize_embedding(embedding):
     buf.seek(0)
     return Image.open(buf)
 
-# --- Main Processing Function ---
+# --- CORRECTED Main Processing Function ---
 def comprehensive_analysis(image, task, architecture, text_prompt, box_threshold, text_threshold):
-    """Performs the main analysis, including dynamic font calculation."""
+    """Performs the main analysis with corrected logic."""
     if image is None:
         raise gr.Error("Please upload an image first!")
 
     # Calculate a dynamic font size based on image width.
-    # The font size will be 4% of the image width, with a minimum size of 15.
     dynamic_font_size = max(15, int(image.width * 0.04))
     try:
         font = ImageFont.truetype("arial.ttf", dynamic_font_size)
     except IOError:
         font = ImageFont.load_default(size=dynamic_font_size)
 
-    if task == "Zero-Shot Detection":
-        architecture = "grounding_dino"
-    
+    # Get a fresh model instance to avoid stateful errors
     model = get_model(task, architecture)
     outputs = {"annotated_image": None, "model_info": "", "classes_info": "", "embedding_plot": None}
 
@@ -126,7 +130,7 @@ def comprehensive_analysis(image, task, architecture, text_prompt, box_threshold
         features = model.extract_features(image)
         outputs["model_info"] = f"Architecture: {architecture.upper()}\nTask: {task}\nDevice: {model.device}"
         outputs["classes_info"] = f"Classes: {model.get_classes()}"
-    else: # Zero-Shot Detection
+    else:  # Zero-Shot Detection
         if not text_prompt:
             raise gr.Error("Please provide a text prompt for Zero-Shot Detection.")
         
@@ -136,17 +140,18 @@ def comprehensive_analysis(image, task, architecture, text_prompt, box_threshold
             box_threshold=box_threshold,
             text_threshold=text_threshold
         )
-        
         outputs["annotated_image"] = draw_dino_predictions(image, results, font=font)
         features = model.extract_features(image, text_prompt=text_prompt)
-        outputs["model_info"] = f"Architecture: {architecture.upper()}\nTask: {task}\nDevice: {model.device}\nHF Model ID: {model.model.config._name_or_path}"
+        outputs["model_info"] = f"Architecture: GROUNDING_DINO\nTask: {task}\nDevice: {model.device}\nHF Model ID: {model.model.config._name_or_path}"
         outputs["classes_info"] = f"Prompt: '{text_prompt}'"
 
+    # Process features for visualization
     if isinstance(features, dict):
         outputs["embedding_plot"] = visualize_embedding(features.get('last_hidden_state'))
     else:
         outputs["embedding_plot"] = visualize_embedding(features)
     
+    # Correctly placed return statement ensures all outputs are always returned
     return outputs["annotated_image"], outputs["model_info"], outputs["classes_info"], outputs["embedding_plot"]
 
 # --- Gradio UI ---