Update app.py

app.py

@@ -137,17 +137,22 @@ class ModelManager:
             base_vision_model_name = self.config.get("base_vision_model")
             print(f"Loading vision model: {base_vision_model_name}")
             
-            # --- UPDATED LOADING LOGIC ---
             is_dinov3_8bit = "dinov3" in base_vision_model_name and "8bit" in base_vision_model_name
             
+            # --- UPDATED LOGIC v5 ---
+            # For 8-bit, the repo contains everything, including the processor.
+            # For others, we load from their base name.
+            processor_name = base_vision_model_name if is_dinov3_8bit else self.config.get("base_vision_model")
+            self.hf_processor = AutoProcessor.from_pretrained(processor_name, trust_remote_code=True) # <-- THE ONLY CHANGE IS HERE
+            
             if is_dinov3_8bit:
-                # Use your 8-bit model from the Hub
-                self.hf_processor = AutoProcessor.from_pretrained("facebook/dinov3-base") # Processor is usually from the base model
                 self.vision_model = AutoModel.from_pretrained(
-                    base_vision_model_name,
-                    load_in_8bit=True,
-                    trust_remote_code=True
+                    base_vision_model_name, load_in_8bit=True, trust_remote_code=True
                 ).eval()
+            else:
+                self.vision_model = AutoModel.from_pretrained(
+                    base_vision_model_name, torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
+                ).to(DEVICE).eval()
             else: # For SigLIP or other non-8bit models
                 self.hf_processor = AutoProcessor.from_pretrained(base_vision_model_name)
                 self.vision_model = AutoModel.from_pretrained(
@@ -239,10 +244,9 @@ def predict_anatomy_v3(image: Image.Image, model_name: str):
         return {"Error": str(e)}
 
 # --- Gradio Interface ---
-# (Unchanged)
 DESCRIPTION = """
 ## Lumi's Anatomy Flaw Classifier Demo ✨
-Select a model from the dropdown, then upload an image to classify its anatomy/structure.
+Select a model from the dropdown, then upload an image to classify its anatomy/structural correctness.
 """
 EXAMPLE_DIR = "examples"
 examples = []
@@ -250,7 +254,20 @@ if os.path.isdir(EXAMPLE_DIR):
     examples = [os.path.join(EXAMPLE_DIR, fname) for fname in sorted(os.listdir(EXAMPLE_DIR)) if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
 
 default_model = list(MODEL_CATALOG.keys())[0]
-interface = gr.Interface(fn=predict_anatomy_v3, inputs=[gr.Image(type="pil", label="Input Image"), gr.Dropdown(choices=list(MODEL_CATALOG.keys()), value=default_model, label="Classifier Model")], outputs=gr.Label(label="Class Probabilities", num_top_classes=2), title="Lumi's Anatomy Classifier", description=DESCRIPTION, examples=examples if examples else None, allow_flagging="never", cache_examples=False)
+
+interface = gr.Interface(
+    fn=predict_anatomy_v3,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Dropdown(choices=list(MODEL_CATALOG.keys()), value=default_model, label="Classifier Model")
+    ],
+    outputs=gr.Label(label="Class Probabilities", num_top_classes=2),
+    title="Lumi's Anatomy Classifier",
+    description=DESCRIPTION,
+    examples=examples if examples else None,
+    allow_flagging="never",
+    cache_examples=True
+)
 
 if __name__ == "__main__":
     try:
@@ -258,4 +275,5 @@ if __name__ == "__main__":
         model_manager.load_model(default_model)
     except Exception as e:
         print(f"WARNING: Could not pre-load default model. Error: {e}")
+        
     interface.launch()