Update app.py

app.py

@@ -110,8 +110,9 @@ MODEL_CATALOG = {
         "config_filename": "AnatomyFlaws-v15.5_dinov3_7b_bnb_fl.config.json",
         "head_filename": "AnatomyFlaws-v15.5_dinov3_7b_bnb_fl_s3K_best_val.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-int4",
+        # "vision_model_repo_id": "Enferlain/dinov3-vit7b16-pretrain-lvd1689m-8bit" bnb 8bit
+        # "vision_model_repo_id": "Enferlain/dinov3-vit7b16-pretrain-lvd1689m-int4", int4 
+        "vision_model_repo_id": "facebook/dinov3-vit7b16-pretrain-lvd1689m",
     },
     "AnatomyFlaws-v14.7 (SigLIP naflex)": {
         "repo_id": "Enferlain/lumi-classifier",
@@ -125,61 +126,89 @@ MODEL_CATALOG = {
 # --- Model Manager Class ---
 class ModelManager:
     def __init__(self, catalog: Dict[str, Dict[str, str]]):
-        self.catalog = catalog; self.current_model_name: str = None; self.vision_model: nn.Module = None
-        self.hf_processor: Any = None; self.head_model: HybridHeadModel = None
-        self.labels: Dict[int, str] = None; self.config: Dict[str, Any] = None
+        self.catalog = catalog
+        self.current_model_name: str = None
+        self.vision_model: nn.Module = None
+        self.hf_processor: Any = None
+        self.head_model: HybridHeadModel = None
+        self.labels: Dict[int, str] = None
+        self.config: Dict[str, Any] = None
 
     def load_model(self, model_name: str):
-        if model_name == self.current_model_name: return
-        if model_name not in self.catalog: raise ValueError(f"Model '{model_name}' not found.")
+        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 = model_info["repo_id"]
         config_filename = model_info["config_filename"]
         head_filename = model_info["head_filename"]
         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)
-            
+            with open(config_path, 'r', encoding='utf-8') as f:
+                self.config = json.load(f)
+
             print(f"Loading vision model: {vision_model_repo_id}")
-            
             self.hf_processor = AutoProcessor.from_pretrained(vision_model_repo_id, trust_remote_code=True)
-            
-            # --- 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.float32,
-                    device_map="cpu", # Force to CPU
-                    trust_remote_code=True
-                ).eval()
-            else: # Standard model loading (for SigLIP or GPU environments)
+
+            # --- UPDATED: CPU-compatible loading logic ---
+            if DEVICE == "cpu":
+                # For CPU, load unquantized model with BF16 (original format)
+                print("Loading unquantized model for CPU...")
+                try:
+                    self.vision_model = AutoModel.from_pretrained(
+                        vision_model_repo_id,
+                        torch_dtype=torch.bfloat16,  # Keep original BF16 format
+                        device_map={"": "cpu"},      # Force CPU device mapping
+                        trust_remote_code=True
+                    ).eval()
+                    print("Successfully loaded model in BF16 format.")
+                except Exception as bf16_error:
+                    print(f"BF16 loading failed: {bf16_error}")
+                    print("Falling back to FP32...")
+                    self.vision_model = AutoModel.from_pretrained(
+                        vision_model_repo_id,
+                        torch_dtype=torch.float32,  # Fallback to FP32
+                        device_map={"": "cpu"},
+                        trust_remote_code=True
+                    ).eval()
+                    print("Successfully loaded model in FP32 format.")
+            else:
+                # For GPU environments (unchanged)
                 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
+
+            # Load classifier head (unchanged)
             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')
             head_params = self.config.get("predictor_params", self.config)
             self.head_model = HybridHeadModel(
-                features=head_params.get("features"), hidden_dim=head_params.get("hidden_dim"),
-                num_classes=self.config.get("num_classes"), use_attention=head_params.get("use_attention"),
-                num_attn_heads=head_params.get("num_attn_heads"), attn_dropout=head_params.get("attn_dropout"),
-                num_res_blocks=head_params.get("num_res_blocks"), dropout_rate=head_params.get("dropout_rate"),
-                output_mode=head_params.get("output_mode", "linear"))
+                features=head_params.get("features"), 
+                hidden_dim=head_params.get("hidden_dim"),
+                num_classes=self.config.get("num_classes"), 
+                use_attention=head_params.get("use_attention"),
+                num_attn_heads=head_params.get("num_attn_heads"), 
+                attn_dropout=head_params.get("attn_dropout"),
+                num_res_blocks=head_params.get("num_res_blocks"), 
+                dropout_rate=head_params.get("dropout_rate"),
+                output_mode=head_params.get("output_mode", "linear")
+            )
             self.head_model.load_state_dict(state_dict, strict=True)
             self.head_model.to(DEVICE).eval()
+
             raw_labels = self.config.get("labels", {'0': 'Bad', '1': 'Good'})
             self.labels = {int(k): (v['name'] if isinstance(v, dict) else v) for k, v in raw_labels.items()}
             self.current_model_name = model_name
             print(f"Successfully loaded '{model_name}'.")
+
         except Exception as e:
             self.current_model_name = None
             raise RuntimeError(f"Failed to load model '{model_name}': {e}\n{traceback.format_exc()}")