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_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') } 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. """ # 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()