Update app.py

app.py

@@ -105,15 +105,19 @@ class HybridHeadModel(nn.Module):
 
 # --- Model Catalog ---
 MODEL_CATALOG = {
-    "AnatomyFlaws-v15.5 (DINOv3 7b 8-bit)": { # <-- Renamed for clarity
+    "AnatomyFlaws-v15.5 (DINOv3 7b 8-bit)": {
         "repo_id": "Enferlain/lumi-classifier",
         "config_filename": "AnatomyFlaws-v15.5_dinov3_7b_bnb_fl.config.json",
-        "head_filename": "AnatomyFlaws-v15.5_dinov3_7b_bnb_fl_s4K.safetensors"
+        "head_filename": "AnatomyFlaws-v15.5_dinov3_7b_bnb_fl_s4K.safetensors",
+        # Explicitly define the vision model repo ID to prevent errors
+        "vision_model_repo_id": "Enferlain/dinov3-vit7b16-pretrain-lvd1689m-8bit"
     },
     "AnatomyFlaws-v14.7 (SigLIP naflex)": {
         "repo_id": "Enferlain/lumi-classifier",
         "config_filename": "AnatomyFlaws-v14.7_adabelief_fl_naflex_4670.config.json",
-        "head_filename": "AnatomyFlaws-v14.7_adabelief_fl_naflex_4670_s2K.safetensors"
+        "head_filename": "AnatomyFlaws-v14.7_adabelief_fl_naflex_4670_s2K.safetensors",
+        # The base SigLIP model is not custom, so we use its official ID
+        "vision_model_repo_id": "google/siglip2-so400m-patch16-naflex"
     },
 }
 
@@ -128,32 +132,32 @@ class ModelManager:
         if model_name == self.current_model_name: return
         if model_name not in self.catalog: raise ValueError(f"Model '{model_name}' not found.")
         print(f"Switching to model: {model_name}...")
+        
         model_info = self.catalog[model_name]
-        repo_id, config_filename, head_filename = model_info["repo_id"], model_info["config_filename"], model_info["head_filename"]
+        repo_id = model_info["repo_id"]
+        config_filename = model_info["config_filename"]
+        head_filename = model_info["head_filename"]
+        # --- NEW: Use the reliable repo ID from our catalog ---
+        vision_model_repo_id = model_info["vision_model_repo_id"]
+        
         try:
             config_path = hf_hub_download(repo_id=repo_id, filename=config_filename)
             with open(config_path, 'r', encoding='utf-8') as f: self.config = json.load(f)
             
-            base_vision_model_name = self.config.get("base_vision_model")
-            print(f"Loading vision model: {base_vision_model_name}")
+            print(f"Loading vision model: {vision_model_repo_id}")
             
-            is_dinov3_8bit = "dinov3" in base_vision_model_name and "8bit" in base_vision_model_name
+            # Load processor and model using our trusted repo ID
+            self.hf_processor = AutoProcessor.from_pretrained(vision_model_repo_id, trust_remote_code=True)
             
-            # --- 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:
-                self.vision_model = AutoModel.from_pretrained(
-                    base_vision_model_name, load_in_8bit=True, trust_remote_code=True
-                ).eval()
+            is_8bit_model = "8bit" in vision_model_repo_id
+            if is_8bit_model:
+                self.vision_model = AutoModel.from_pretrained(vision_model_repo_id, 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
+                    vision_model_repo_id, torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
                 ).to(DEVICE).eval()
             
+            # The rest of the function continues as before
             head_model_path = hf_hub_download(repo_id=repo_id, filename=head_filename)
             print(f"Loading head model: {head_filename}")
             state_dict = load_file(head_model_path, device='cpu')