Create app.py

app.py

@@ -0,0 +1,286 @@
+import os
+import json
+import traceback
+from typing import Optional, Tuple, Union, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image, PngImagePlugin
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+from transformers import AutoProcessor, AutoModel, AutoImageProcessor
+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
+
+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) ---
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+    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.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or int(in_features * 8 / 3 / 2 * 2 )
+        hidden_features = (hidden_features + 1) // 2 * 2
+        self.w12 = nn.Linear(in_features, hidden_features * 2, bias=False)
+        self.act = act_layer()
+        self.dropout1 = nn.Dropout(dropout)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=False)
+        self.dropout2 = nn.Dropout(dropout)
+    def forward(self, x):
+        gate_val, up_val = self.w12(x).chunk(2, dim=-1)
+        x = self.dropout1(self.act(gate_val) * up_val)
+        x = self.dropout2(self.w3(x))
+        return x
+
+class ResBlockRMS(nn.Module):
+    def __init__(self, ch: int, dropout: float = 0.0, rms_norm_eps: float = 1e-6):
+        super().__init__()
+        self.norm = RMSNorm(ch, eps=rms_norm_eps)
+        self.ffn = SwiGLUFFN(in_features=ch, dropout=dropout)
+    def forward(self, x):
+        return x + self.ffn(self.norm(x))
+
+class HybridHeadModel(nn.Module):
+    def __init__(self, features: int, hidden_dim: int = 1280, num_classes: int = 2, use_attention: bool = True,
+                 num_attn_heads: int = 16, attn_dropout: float = 0.1, num_res_blocks: int = 3,
+                 dropout_rate: float = 0.1, rms_norm_eps: float = 1e-6, output_mode: str = 'linear'):
+        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:
+            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
+                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}")
+
+            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))
+        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
+        mlp_layers.append(SwiGLUFFN(hidden_dim, hidden_features=down_proj_hidden, out_features=down_proj_hidden, dropout=dropout_rate))
+        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}'.")
+        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_s9K.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')
+}
+print(f"Using Labels: {LABEL_NAMES}")
+
+# --- Prediction Function ---
+def predict_anatomy(image: Image.Image):
+    """Takes PIL Image, returns dict of class probabilities."""
+    if image is None: return {"Error": "No image provided"}
+    try:
+        pil_image = image.convert("RGB")
+
+        # 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
+            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
+        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()
+        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()}
+        return output_probs
+
+    except Exception as e:
+        print(f"Error during prediction: {e}\n{traceback.format_exc()}")
+        return {"Error": str(e)}
+
+# --- Gradio Interface ---
+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.
+Model Checkpoint: **AnatomyFlaws-v11.3_..._s9K** (or specify which one).
+"""
+
+# 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
+)
+
+if __name__ == "__main__":
+    interface.launch()