app.py

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()