Update app.py

app.py

@@ -1,31 +1,33 @@
 import os
 import json
 import traceback
-from typing import Optional, Tuple, Union, List
+from typing import Dict, Any
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from PIL import Image, PngImagePlugin
+from PIL import Image
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
-from transformers import AutoProcessor, AutoModel, AutoImageProcessor
+from transformers import AutoProcessor, AutoModel
 import gradio as gr
-import math # Added math
 
 # --- Device Setup ---
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# Use float16 for vision model on CUDA for speed/memory, but head expects float32
-VISION_DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
-HEAD_DTYPE = torch.float32 # Head usually trained/stable in float32
+# For 8-bit models, the vision dtype is handled by bitsandbytes
+# We still need HEAD_DTYPE for our classifier head
+HEAD_DTYPE = torch.float32
+
+# --- DINOv3 Specific Constants ---
+DINOV3_PATCH_SIZE = 16
+MAX_DINOV3_RESOLUTION = 4096
 
 print(f"Using device: {DEVICE}")
-print(f"Vision model dtype: {VISION_DTYPE}")
 print(f"Head model dtype: {HEAD_DTYPE}")
 
 
 # --- Model Definitions (Copied from hybrid_model.py) ---
-
+# (RMSNorm, SwiGLUFFN, ResBlockRMS, HybridHeadModel classes are unchanged and go here)
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
@@ -36,8 +38,6 @@ class RMSNorm(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         output = self._norm(x.float()).type_as(x)
         return output * self.weight
-    def extra_repr(self) -> str:
-        return f"{tuple(self.weight.shape)}, eps={self.eps}"
 
 class SwiGLUFFN(nn.Module):
     def __init__(self, in_features: int, hidden_features: int = None, out_features: int = None, act_layer: nn.Module = nn.SiLU, dropout: float = 0.):
@@ -71,22 +71,17 @@ class HybridHeadModel(nn.Module):
         super().__init__()
         self.features = features; self.hidden_dim = hidden_dim; self.num_classes = num_classes
         self.use_attention = use_attention; self.output_mode = output_mode.lower()
-        # --- Optional Self-Attention Layer ---
         self.attention = None; self.norm_attn = None
         if self.use_attention:
-            actual_num_heads = num_attn_heads # Adjust head logic needed here if features != 1152
-            # Simple head adjustment:
+            actual_num_heads = num_attn_heads
             if features % num_attn_heads != 0:
-                possible_heads = [h for h in [1, 2, 4, 8, 16] if features % h == 0]
-                if not possible_heads: actual_num_heads = 1 # Fallback to 1 head if no divisors found
+                possible_heads = [h for h in [1, 2, 4, 8, 16, 32] if features % h == 0] # Expanded list
+                if not possible_heads: actual_num_heads = 1
                 else: actual_num_heads = min(possible_heads, key=lambda x: abs(x-num_attn_heads))
-                if actual_num_heads != num_attn_heads: print(f"HybridHead Warning: Adjusting heads {num_attn_heads}->{actual_num_heads}")
-
+                if actual_num_heads != num_attn_heads: print(f"HybridHead Warning: Adjusting heads {num_attn_heads}->{actual_num_heads} for features={features}")
             self.attention = nn.MultiheadAttention(features, actual_num_heads, dropout=attn_dropout, batch_first=True, bias=True)
             self.norm_attn = RMSNorm(features, eps=rms_norm_eps)
-        # --- MLP Head ---
-        mlp_layers = []
-        mlp_layers.append(nn.Linear(features, hidden_dim)); mlp_layers.append(RMSNorm(hidden_dim, eps=rms_norm_eps))
+        mlp_layers = [nn.Linear(features, hidden_dim), RMSNorm(hidden_dim, eps=rms_norm_eps)]
         for _ in range(num_res_blocks): mlp_layers.append(ResBlockRMS(hidden_dim, dropout=dropout_rate, rms_norm_eps=rms_norm_eps))
         mlp_layers.append(RMSNorm(hidden_dim, eps=rms_norm_eps))
         down_proj_hidden = hidden_dim // 2
@@ -94,192 +89,173 @@ class HybridHeadModel(nn.Module):
         mlp_layers.append(RMSNorm(down_proj_hidden, eps=rms_norm_eps))
         mlp_layers.append(nn.Linear(down_proj_hidden, num_classes))
         self.mlp_head = nn.Sequential(*mlp_layers)
-        # --- Validate Output Mode ---
-        # (Warnings can be added here if desired, but functionality handled in forward)
 
     def forward(self, x: torch.Tensor):
         if self.use_attention and self.attention is not None:
             x_seq = x.unsqueeze(1); attn_output, _ = self.attention(x_seq, x_seq, x_seq); x = self.norm_attn(x + attn_output.squeeze(1))
-        logits = self.mlp_head(x.to(HEAD_DTYPE)) # Ensure input to MLP has correct dtype
-        # --- Apply Final Activation ---
-        output = None
-        if self.output_mode == 'linear': output = logits
-        elif self.output_mode == 'sigmoid': output = torch.sigmoid(logits)
-        elif self.output_mode == 'softmax': output = F.softmax(logits, dim=-1)
-        elif self.output_mode == 'tanh_scaled': output = (torch.tanh(logits) + 1.0) / 2.0
-        else: raise RuntimeError(f"Invalid output_mode '{self.output_mode}'.")
+        logits = self.mlp_head(x.to(HEAD_DTYPE))
+        output_mode = self.output_mode
+        if output_mode == 'linear': output = logits
+        elif output_mode == 'sigmoid': output = torch.sigmoid(logits)
+        elif output_mode == 'softmax': output = F.softmax(logits, dim=-1)
+        elif output_mode == 'tanh_scaled': output = (torch.tanh(logits) + 1.0) / 2.0
+        else: raise RuntimeError(f"Invalid output_mode '{output_mode}'.")
         if self.num_classes == 1 and output.ndim == 2 and output.shape[1] == 1: output = output.squeeze(-1)
         return output
 
-# --- Constants and Model Loading ---
-
-# Option 1: Files are in the Space repo (e.g., in a 'model' folder)
-# MODEL_DIR = "model"
-# HEAD_MODEL_FILENAME = "AnatomyFlaws-v11.3_adabelief_fl_naflex_3000_s9K.safetensors"
-# CONFIG_FILENAME = "AnatomyFlaws-v11.3_adabelief_fl_naflex_3000.config.json" # Assuming config matches base name
-# HEAD_MODEL_PATH = os.path.join(MODEL_DIR, HEAD_MODEL_FILENAME)
-# CONFIG_PATH = os.path.join(MODEL_DIR, CONFIG_FILENAME)
-
-# Option 2: Download from Hub
-# Replace with your HF username and repo name
-HUB_REPO_ID = "Enferlain/lumi-classifier" # Or wherever you uploaded the model
-# Use the specific checkpoint you want (e.g., s9k or the best_val one)
-HEAD_MODEL_FILENAME = "AnatomyFlaws-v11.3_adabelief_fl_naflex_3000_s6K_best_val.safetensors"
-# Usually config corresponds to the base run name, not a specific step
-CONFIG_FILENAME = "AnatomyFlaws-v11.3_adabelief_fl_naflex_3000.config.json"
-
-print("Downloading model files if necessary...")
-try:
-    HEAD_MODEL_PATH = hf_hub_download(repo_id=HUB_REPO_ID, filename=HEAD_MODEL_FILENAME)
-    CONFIG_PATH = hf_hub_download(repo_id=HUB_REPO_ID, filename=CONFIG_FILENAME)
-    print("Files downloaded/found successfully.")
-except Exception as e:
-    print(f"ERROR downloading files from {HUB_REPO_ID}: {e}")
-    print("Please ensure the files exist on the Hub or place them in a local 'model' folder.")
-    # Optionally exit or fallback
-    exit(1) # Exit if essential files aren't available
-
-
-# --- Load Config ---
-print(f"Loading config from: {CONFIG_PATH}")
-config = {}
-try:
-    with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
-        config = json.load(f)
-except Exception as e:
-    print(f"ERROR loading config file: {e}"); exit(1)
-
-# --- Load Vision Model ---
-BASE_VISION_MODEL_NAME = config.get("base_vision_model", "google/siglip2-so400m-patch16-naflex")
-print(f"Loading vision model: {BASE_VISION_MODEL_NAME}")
-try:
-    hf_processor = AutoProcessor.from_pretrained(BASE_VISION_MODEL_NAME)
-    vision_model = AutoModel.from_pretrained(
-        BASE_VISION_MODEL_NAME, torch_dtype=VISION_DTYPE
-    ).to(DEVICE).eval()
-    print("Vision model loaded.")
-except Exception as e:
-    print(f"ERROR loading vision model: {e}"); exit(1)
-
-# --- Load HybridHeadModel ---
-print(f"Loading head model: {HEAD_MODEL_PATH}")
-head_model = None
-try:
-    state_dict = load_file(HEAD_MODEL_PATH, device='cpu')
-    # Infer details from config - use defaults matching the successful run
-    features = config.get("features", 1152)
-    num_classes = config.get("num_classes", 2) # Should be 2 for focal loss run
-    output_mode = config.get("output_mode", "linear") # Should be linear
-    hidden_dim = config.get("hidden_dim", 1280)
-    num_res_blocks = config.get("num_res_blocks", 3)
-    dropout_rate = config.get("dropout_rate", 0.3) # Use the high dropout from best run
-    use_attention = config.get("use_attention", True) # Use attention was likely True
-    num_attn_heads = config.get("num_attn_heads", 16)
-    attn_dropout = config.get("attn_dropout", 0.3) # Use the high dropout
-    rms_norm_eps= config.get("rms_norm_eps", 1e-6)
-
-    head_model = HybridHeadModel(
-        features=features, hidden_dim=hidden_dim, num_classes=num_classes,
-        use_attention=use_attention, num_attn_heads=num_attn_heads, attn_dropout=attn_dropout,
-        num_res_blocks=num_res_blocks, dropout_rate=dropout_rate, rms_norm_eps=rms_norm_eps,
-        output_mode=output_mode
-    )
-    missing, unexpected = head_model.load_state_dict(state_dict, strict=False)
-    if missing: print(f"Warning: Missing keys loading head: {missing}")
-    if unexpected: print(f"Warning: Unexpected keys loading head: {unexpected}")
-    head_model.to(DEVICE).eval()
-    print("Head model loaded.")
-except Exception as e:
-    print(f"ERROR loading head model: {e}"); exit(1)
-
-# --- Label Mapping ---
-# Assume labels are '0': Bad, '1': Good from config or default
-LABELS = config.get("labels", {'0': 'Bad Anatomy', '1': 'Good Anatomy'})
-LABEL_NAMES = {
-    0: LABELS.get('0', 'Class 0'),
-    1: LABELS.get('1', 'Class 1')
+# --- Model Catalog ---
+MODEL_CATALOG = {
+    "AnatomyFlaws-v15.5 (DINOv3 7b 8-bit)": { # <-- Renamed for clarity
+        "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"
+    },
+    "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"
+    },
 }
-print(f"Using Labels: {LABEL_NAMES}")
 
-# --- Prediction Function ---
-def predict_anatomy(image: Image.Image):
-    """Takes PIL Image, returns dict of class probabilities."""
+# --- 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
+
+    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.")
+        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"]
+        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}")
+            
+            # --- UPDATED LOADING LOGIC ---
+            is_dinov3_8bit = "dinov3" in base_vision_model_name and "8bit" in base_vision_model_name
+            
+            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
+                ).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(
+                    base_vision_model_name,
+                    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32 # Use a dynamic dtype
+                ).to(DEVICE).eval()
+            
+            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"))
+            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()}")
+
+# --- Global Model Manager Instance ---
+model_manager = ModelManager(MODEL_CATALOG)
+
+# --- Prediction Function (v3 from before) ---
+def predict_anatomy_v3(image: Image.Image, model_name: str):
     if image is None: return {"Error": "No image provided"}
     try:
+        model_manager.load_model(model_name)
         pil_image = image.convert("RGB")
+        emb = None
 
-        # 1. Extract SigLIP NaFlex Embedding
         with torch.no_grad():
-            inputs = hf_processor(images=[pil_image], return_tensors="pt", max_num_patches=1024)
-            pixel_values = inputs.get("pixel_values").to(device=DEVICE, dtype=VISION_DTYPE)
-            attention_mask = inputs.get("pixel_attention_mask").to(device=DEVICE)
-            spatial_shapes = inputs.get("spatial_shapes")
-            model_call_kwargs = {"pixel_values": pixel_values, "attention_mask": attention_mask,
-                                 "spatial_shapes": torch.tensor(spatial_shapes, dtype=torch.long).to(DEVICE)}
-
-            vision_model_component = getattr(vision_model, 'vision_model', vision_model) # Handle potential nesting
-            emb = vision_model_component(**model_call_kwargs).pooler_output
+            base_model_type = model_manager.config.get("base_vision_model", "")
+            if "dinov3" in base_model_type.lower():
+                current_w, current_h = pil_image.size
+                img_to_process = pil_image
+                if max(current_w, current_h) > MAX_DINOV3_RESOLUTION:
+                    scale = MAX_DINOV3_RESOLUTION / max(current_w, current_h)
+                    current_w, current_h = int(current_w * scale), int(current_h * scale)
+                    img_to_process = pil_image.resize((current_w, current_h), Image.Resampling.LANCZOS)
+                new_w = ((current_w + DINOV3_PATCH_SIZE - 1) // DINOV3_PATCH_SIZE) * DINOV3_PATCH_SIZE
+                new_h = ((current_h + DINOV3_PATCH_SIZE - 1) // DINOV3_PATCH_SIZE) * DINOV3_PATCH_SIZE
+                if new_w != current_w or new_h != current_h:
+                    img_to_process = img_to_process.resize((new_w, new_h), Image.Resampling.LANCZOS)
+                inputs = model_manager.hf_processor(images=[img_to_process], return_tensors="pt")
+                # For 8-bit, send inputs to the same device as the model
+                pixel_values = inputs.pixel_values.to(model_manager.vision_model.device)
+                outputs = model_manager.vision_model(pixel_values=pixel_values)
+                last_hidden_state = outputs.last_hidden_state
+                nreg = getattr(model_manager.vision_model.config, 'num_register_tokens', 0)
+                patch_embeddings = last_hidden_state[:, 1 + nreg:]
+                emb = torch.mean(patch_embeddings, dim=1)
+            elif "siglip" in base_model_type.lower():
+                inputs = model_manager.hf_processor(images=[pil_image], return_tensors="pt")
+                pixel_values = inputs.get("pixel_values").to(device=DEVICE, dtype=torch.float16)
+                if "naflex" in base_model_type.lower():
+                    attention_mask = inputs.get("pixel_attention_mask").to(device=DEVICE)
+                    spatial_shapes = inputs.get("spatial_shapes")
+                    model_call_kwargs = {"pixel_values": pixel_values, "attention_mask": attention_mask,
+                                         "spatial_shapes": torch.tensor(spatial_shapes, dtype=torch.long).to(DEVICE)}
+                    vision_model_component = getattr(model_manager.vision_model, 'vision_model', model_manager.vision_model)
+                    emb = vision_model_component(**model_call_kwargs).pooler_output
+                else: emb = model_manager.vision_model.get_image_features(pixel_values=pixel_values)
+            else: raise ValueError(f"Unknown base model type for embedding: {base_model_type}")
             if emb is None: raise ValueError("Failed to get embedding.")
-
-            # L2 Norm
             norm = torch.linalg.norm(emb.float(), dim=-1, keepdim=True).clamp(min=1e-8)
             emb_normalized = emb / norm.to(emb.dtype)
-
-        # 2. Obtain Prediction from HybridHeadModel Head
         with torch.no_grad():
-            prediction = head_model(emb_normalized.to(DEVICE, dtype=HEAD_DTYPE))
-
-        # 3. Format Output Probabilities
+            prediction = model_manager.head_model(emb_normalized.to(DEVICE, dtype=HEAD_DTYPE))
         output_probs = {}
-        output_mode = getattr(head_model, 'output_mode', 'linear')
-
-        if head_model.num_classes == 1:
-            logit = prediction.squeeze().item()
-            prob_good = torch.sigmoid(torch.tensor(logit)).item() if output_mode == 'linear' else logit
-            output_probs[LABEL_NAMES[0]] = 1.0 - prob_good
-            output_probs[LABEL_NAMES[1]] = prob_good
-        elif head_model.num_classes == 2:
-            if output_mode == 'linear':
-                probs = F.softmax(prediction.squeeze().float(), dim=-1) # Use float for softmax stability
-            else: # Assume sigmoid or already softmax
-                probs = prediction.squeeze().float()
-            output_probs[LABEL_NAMES[0]] = probs[0].item()
-            output_probs[LABEL_NAMES[1]] = probs[1].item()
+        if model_manager.head_model.num_classes == 2:
+            probs = F.softmax(prediction.squeeze().float(), dim=-1)
+            output_probs[model_manager.labels[0]] = probs[0].item()
+            output_probs[model_manager.labels[1]] = probs[1].item()
         else:
-             output_probs["Error"] = f"Unsupported num_classes: {head_model.num_classes}"
-
-        # Convert to percentage strings for gr.Label maybe? Or keep floats? Keep floats.
-        # output_formatted = {k: f"{v:.1%}" for k, v in output_probs.items()}
+            prob_good = torch.sigmoid(prediction.squeeze()).item()
+            output_probs[model_manager.labels[0]] = 1.0 - prob_good
+            output_probs[model_manager.labels[1]] = prob_good
         return output_probs
-
     except Exception as e:
         print(f"Error during prediction: {e}\n{traceback.format_exc()}")
         return {"Error": str(e)}
 
 # --- Gradio Interface ---
+# (Unchanged)
 DESCRIPTION = """
-## Anatomy Flaw Classifier Demo ✨ (Based on SigLIP Naflex + Hybrid Head)
-Upload an image to classify its anatomy as 'Good' or 'Bad'.
-This model uses embeddings from **google/siglip2-so400m-patch16-naflex**
-and a custom **HybridHeadModel** fine-tuned for anatomy classification.
+## Lumi's Anatomy Flaw Classifier Demo ✨
+Select a model from the dropdown, then upload an image to classify its anatomy/structure.
 """
-
-# Add example images if you have some in an 'examples' folder in the Space repo
 EXAMPLE_DIR = "examples"
 examples = []
 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'))]
 
-interface = gr.Interface(
-    fn=predict_anatomy,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Label(label="Class Probabilities", num_top_classes=2), # Show top 2 classes
-    title="Lumi's Anatomy Classifier Demo",
-    description=DESCRIPTION,
-    examples=examples if examples else None,
-    allow_flagging="never",
-    cache_examples=False # Disable caching if examples change or loading is fast
-)
+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)
 
 if __name__ == "__main__":
+    try:
+        print("Pre-loading default model...")
+        model_manager.load_model(default_model)
+    except Exception as e:
+        print(f"WARNING: Could not pre-load default model. Error: {e}")
     interface.launch()