Update app.py

app.py

@@ -110,7 +110,8 @@ MODEL_CATALOG = {
         "config_filename": "AnatomyFlaws-v15.5_dinov3_7b_bnb_fl.config.json",
         "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"
+        # "vision_model_repo_id": "Enferlain/dinov3-vit7b16-pretrain-lvd1689m-8bit"
+        "vision_model_repo_id": "Enferlain/dinov3-vit7b16-pretrain-lvd1689m-int4",
     },
     "AnatomyFlaws-v14.7 (SigLIP naflex)": {
         "repo_id": "Enferlain/lumi-classifier",
@@ -137,7 +138,6 @@ class ModelManager:
         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:
@@ -146,15 +146,21 @@ class ModelManager:
             
             print(f"Loading vision model: {vision_model_repo_id}")
             
-            # Load processor and model using our trusted repo ID
             self.hf_processor = AutoProcessor.from_pretrained(vision_model_repo_id, trust_remote_code=True)
             
-            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:
+            # --- NEW: Correct loading logic for INT4 vs. standard models ---
+            if "int4" in vision_model_repo_id.lower():
+                print("INT4 model detected. Loading for CPU.")
                 self.vision_model = AutoModel.from_pretrained(
-                    vision_model_repo_id, torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32
+                    vision_model_repo_id,
+                    torch_dtype=torch.float32,
+                    device_map="cpu", # Force to CPU
+                    trust_remote_code=True
+                ).eval()
+            else: # Standard model loading (for SigLIP or GPU environments)
+                self.vision_model = AutoModel.from_pretrained(
+                    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