utils/complete_model.py

import os
from pathlib import Path
import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, GPT2Tokenizer

from utils.modifiedGPT2 import create_decoder

from utils.layer_mask import gaussian_layer_stack_pipeline

_REPO_ROOT = Path(os.environ.get("LANA_ARXIV_REPO_ROOT", Path(__file__).resolve().parents[1]))
_BUNDLED_VISION_MODEL = (_REPO_ROOT / "bundled_backbones" / "vision_encoder").resolve().as_posix()
_BUNDLED_SEGMENTATION_MODEL = (_REPO_ROOT / "bundled_backbones" / "segmenter_encoder").resolve().as_posix()
_BUNDLED_TEXT_MODEL = (_REPO_ROOT / "bundled_backbones" / "text_decoder").resolve().as_posix()

class DINOEncoder(nn.Module):
    def __init__(self, model_id=_BUNDLED_VISION_MODEL, freeze=True):
        super().__init__()
        self.model = AutoModel.from_config(AutoConfig.from_pretrained(model_id, local_files_only=True))
        if freeze:
            for p in self.model.parameters():
                p.requires_grad = False
    @torch.no_grad()
    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """
        pixel_values: [B, C, H, W]
        returns patches: [B, Np, Cenc]
        """
        out = self.model(pixel_values=pixel_values)
        tokens = out.last_hidden_state  # [B, 1+Np, Cenc] (CLS + patches) for ViT-like
        # Skip a few special tokens if your backbone adds them; adjust as needed.
        patches = tokens[:, 5:, :]  # [B, Np, Cenc]
        return patches

class DinoUNet(nn.Module):
    def __init__(self, model_name=_BUNDLED_SEGMENTATION_MODEL, freeze=True):
        super().__init__()
        self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(model_name, local_files_only=True))
        # NOTE: confirm channels of the chosen hidden state; 768 is common for small convnext/dinov3
        self.channel_adapter = nn.Conv2d(768, 512, kernel_size=1)
        self.decoder = nn.Sequential(
            nn.Conv2d(512, 256, 3, padding=1), nn.ReLU(inplace=True),
            nn.ConvTranspose2d(256, 128, 2, stride=2), nn.ReLU(inplace=True),
            nn.ConvTranspose2d(128, 64, 2, stride=2), nn.ReLU(inplace=True),
            nn.Conv2d(64, 1, 1)
        )
        if freeze:
            for m in (self.encoder, self.channel_adapter, self.decoder):
                for p in m.parameters():
                    p.requires_grad = False
    
    @torch.no_grad()
    def forward(self, x: torch.Tensor, num_layers: int) -> torch.Tensor:
        """
        x: [B, C, H, W]; returns mask: [B, 1, H', W'] (your upsampling stack defines H',W')
        """
        enc_feats = self.encoder(x, output_hidden_states=True, return_dict=True)
        # take the last 4D feature map from hidden_states
        feats = next(h for h in reversed(enc_feats.hidden_states) if isinstance(h, torch.Tensor) and h.ndim == 4)
        feats = self.channel_adapter(feats)
        pred = self.decoder(feats)                    # (B,1,h,w)
        _, _, segmentation_mask = gaussian_layer_stack_pipeline(pred, n_layers = num_layers)
        return segmentation_mask    # [B, num_layers, h, w]


class LinearProjection(nn.Module):
    def __init__(self, input_dim=384, output_dim=768, freeze=False):
        super().__init__()
        self.proj = nn.Linear(input_dim, output_dim)
        if freeze:
            for p in self.proj.parameters():
                p.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, Np, input_dim] -> [B, Np, output_dim]
        return self.proj(x)


class CustomModel(nn.Module):
    def __init__(
        self,
        device: str = "cuda",
        ENCODER_MODEL_PATH: str | None = "dino_encoder.pth",
        SEGMENTER_MODEL_PATH: str | None = "dino_segmenter.pth",
        DECODER_MODEL_PATH: str | None = "dino_decoder.pth",
        LINEAR_PROJECTION_PATH: str | None = "linear_projection.pth",
        freeze_encoder: bool = True,
        freeze_segmenter: bool = True,
        freeze_linear_projection: bool = False,
        freeze_decoder: bool = False,
        attention_implementation: str = "sdpa",
    ):
        super().__init__()
        self.device = torch.device(device)

        # Encoder
        self.encoder = DINOEncoder()
        if ENCODER_MODEL_PATH and os.path.exists(ENCODER_MODEL_PATH):
            self.encoder.load_state_dict(torch.load(ENCODER_MODEL_PATH, map_location="cpu"), strict=False)
            print("Loaded encoder weights from", ENCODER_MODEL_PATH)
        if freeze_encoder:
            self.encoder.eval()

        # Segmenter
        self.segmenter = DinoUNet()
        if SEGMENTER_MODEL_PATH and os.path.exists(SEGMENTER_MODEL_PATH):
            self.segmenter.load_state_dict(torch.load(SEGMENTER_MODEL_PATH, map_location="cpu"), strict=False)
            print("Loaded segmenter weights from", SEGMENTER_MODEL_PATH)
        if freeze_segmenter:
            self.segmenter.eval()

        # Decoder (modified GPT-2)
        self.decoder = create_decoder(attention=attention_implementation)  # must expose .config.hidden_size & .config.num_hidden_layers
        if DECODER_MODEL_PATH and os.path.exists(DECODER_MODEL_PATH):
            self.decoder.load_state_dict(torch.load(DECODER_MODEL_PATH, map_location="cpu"), strict=False)
            print("Loaded decoder weights from", DECODER_MODEL_PATH)
        if freeze_decoder:
            self.decoder.eval()

        # Linear projection: DINO hidden -> GPT2 hidden
        enc_h = self.encoder.model.config.hidden_size
        dec_h = self.decoder.config.hidden_size
        self.linear_projection = LinearProjection(input_dim=enc_h, output_dim=dec_h)
        if LINEAR_PROJECTION_PATH and os.path.exists(LINEAR_PROJECTION_PATH):
            self.linear_projection.load_state_dict(torch.load(LINEAR_PROJECTION_PATH, map_location="cpu"), strict=False)
            print("Loaded linear projection weights from", LINEAR_PROJECTION_PATH)
        if freeze_linear_projection:
            self.linear_projection.eval()

        # Tokenizer (pad token for GPT-2)
        self.tokenizer = GPT2Tokenizer.from_pretrained(_REPO_ROOT.as_posix(), local_files_only=True)
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.pad_token_id = self.tokenizer.pad_token_id  # ✅ use ID, not string

        self.num_layers = self.decoder.config.num_hidden_layers

        # move everything once
        self.to(self.device)

    def forward(self, pixel_values: torch.Tensor, tgt_ids: torch.Tensor | None = None, **kwargs) -> dict:
        """
        pixel_values: [B,C,H,W], float
        tgt_ids: [B,T], long (token IDs), padded with pad_token_id if any padding is present
        """
        pixel_values = pixel_values.to(self.device, non_blocking=True)

        # Visual path
        patches = self.encoder(pixel_values)                           # [B,Np,Cenc]
        projected_patches = self.linear_projection(patches)            # [B,Np,n_embd]

        # Segmentation path per layer
        segmented_layers = self.segmenter(pixel_values, self.num_layers) # [B,n_layers,H,W] (per current decoder)

        # Text path (optional teacher-forced training)
        labels = None
        if tgt_ids is not None:
            if tgt_ids.dtype != torch.long:
                tgt_ids = tgt_ids.long()
            tgt_ids = tgt_ids.to(self.device, non_blocking=True)       # [B,T]
            text_embeds = self.decoder.transformer.wte(tgt_ids)        # [B,T,n_embd]
            inputs_embeds = torch.cat([projected_patches, text_embeds], dim=1)  # [B,Np+T,n_embd]

            # Labels: ignore prefix tokens (vision) and PADs in text
            B, Np, _ = projected_patches.shape
            labels_prefix = torch.full((B, Np), -100, device=self.device, dtype=torch.long)
            text_labels = tgt_ids.clone()
            text_labels[text_labels == self.pad_token_id] = -100       # ✅ compare to ID
            labels = torch.cat([labels_prefix, text_labels], dim=1)    # [B,Np+T]
        else:
            inputs_embeds = projected_patches

        # Decoder forward
        out = self.decoder(inputs_embeds=inputs_embeds, segmentation_mask=segmented_layers, labels=labels, **kwargs)
        return out
    
    @torch.inference_mode()
    def generate(
        self,
        pixel_values: torch.Tensor,
        max_new_tokens: int = 100,
        output_attentions: bool = False,
    ) -> torch.Tensor:
        """
        pixel_values: [B,C,H,W], float
        returns generated_ids: [B, T]
        """
        pixel_values = pixel_values.to(self.device, non_blocking=True)

        # Visual path
        patches = self.encoder(pixel_values)                           # [B,Np,Cenc]
        projected_patches = self.linear_projection(patches)            # [B,Np,n_embd]

        # Segmentation path per layer
        segmented_layers = self.segmenter(pixel_values, self.num_layers) # [B,n_layers,H,W] (per current decoder)

        # Generate
        output = self.decoder.generate(
            inputs_embeds=projected_patches,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=1.2,
            eos_token_id=self.tokenizer.eos_token_id,
            pad_token_id=self.pad_token_id,
            use_cache=True,
            segmentation_mask=segmented_layers,
            prefix_allowed_length=0,
            plot_attention_mask=False,
            plot_attention_mask_layer=[],
            plot_attention_map=False,
            plot_attention_map_layer=[],
            plot_attention_map_generation=0,
            output_attentions=output_attentions,
            return_dict_in_generate=True,
        )
        # Remove prefix tokens (vision)
        generated_ids = output.sequences#[:, projected_patches.shape[1]:]   # [B,T]
        generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
        return generated_ids, generated_text, output.attentions if output_attentions else None

def create_complete_model(device: str = "cuda", **kwargs) -> CustomModel:
    model = CustomModel(device=device, **kwargs)
    return model

def save_complete_model(model: CustomModel, save_path: str, device: str = "cuda") -> None:
    # Ensure folder exists
    os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)

    # Save on CPU to keep checkpoint portable
    orig_device = next(model.parameters()).device
    model.to("cpu")
    torch.save(model.state_dict(), save_path)
    print(f"Saved complete model weights to {save_path}")

    # Restore model device
    model.to(device if isinstance(device, str) else orig_device)

def save_checkpoint(model: CustomModel, optimizer: torch.optim.Optimizer, save_path: str) -> None:
    # Ensure folder exists
    os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)

    checkpoint = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
    }
    torch.save(checkpoint, save_path)
    print(f"Saved checkpoint to {save_path}")

def load_complete_model(model: CustomModel, load_path: str, device: str = "cpu", strict: bool = True) -> CustomModel:
    if not os.path.exists(load_path):
        print(f"No weights found at {load_path}")
        model.to(device)
        return model

    # Load to CPU first, then move to target device
    state = torch.load(load_path, map_location="cpu")
    missing, unexpected = model.load_state_dict(state, strict=strict)
    if not strict:
        if missing:
            print(f"[load warning] Missing keys: {missing}")
        if unexpected:
            print(f"[load warning] Unexpected keys: {unexpected}")

    model.to(device)
    print(f"Loaded complete model weights from {load_path}")
    return model

def load_checkpoint(model: CustomModel, optimizer: torch.optim.Optimizer, load_path: str, device: str = "cpu") -> tuple[CustomModel, torch.optim.Optimizer]:
    if not os.path.exists(load_path):
        print(f"No checkpoint found at {load_path}")
        model.to(device)
        return model, optimizer

    # Load to CPU first, then move to target device
    checkpoint = torch.load(load_path, map_location="cpu")
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

    model.to(device)
    print(f"Loaded checkpoint from {load_path}")
    return model, optimizer

from transformers import AutoImageProcessor
from PIL import Image
import logging
import re

# Configure basic logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ==============================================================================
# 1. Architecture Definition (MLP)
# ==============================================================================
class EmbeddingClassifier(nn.Module):
    """
    Flexible MLP Classifier: Input Embeddings -> Hidden Layers -> Logits.
    """
    def __init__(self, embedding_dim, num_classes, custom_dims=(512, 256, 256),
                 activation="gelu", dropout=0.05, bn=False, use_layernorm=True):
        super().__init__()
        layers = []
        
        # First layer: Embeddings -> First hidden dimension
        layers.append(nn.Linear(embedding_dim, custom_dims[0]))
        if use_layernorm: layers.append(nn.LayerNorm(custom_dims[0]))
        elif bn: layers.append(nn.BatchNorm1d(custom_dims[0]))
        layers.append(nn.GELU() if activation.lower() == "gelu" else nn.ReLU())
        if dropout > 0: layers.append(nn.Dropout(dropout))

        # Intermediate layers
        for i in range(len(custom_dims) - 1):
            layers.append(nn.Linear(custom_dims[i], custom_dims[i + 1]))
            if use_layernorm: layers.append(nn.LayerNorm(custom_dims[i + 1]))
            elif bn: layers.append(nn.BatchNorm1d(custom_dims[i + 1]))
            layers.append(nn.GELU() if activation.lower() == "gelu" else nn.ReLU())
            if dropout > 0: layers.append(nn.Dropout(dropout))

        # Final layer: Last hidden dim -> Num classes (Logits)
        layers.append(nn.Linear(custom_dims[-1], num_classes))
        self.classifier = nn.Sequential(*layers)

    def forward(self, embeddings):
        return self.classifier(embeddings)


# ==============================================================================
# 2. Prediction Wrapper Class
# ==============================================================================
class ChestXrayPredictor:
    """
    Wrapper class responsible for receiving an image, processing it, 
    and returning class probabilities.
    """
    
    def __init__(self, base_model, classifier, processor, label_cols, device):
        self.base_model = base_model
        self.classifier = classifier
        self.processor = processor
        self.label_cols = label_cols
        self.device = device
        
        # Ensure models are in eval mode
        self.base_model.eval()
        self.classifier.eval()

    def predict(self, image_source):
        """
        Runs inference on a single image.

        Args:
            image_source: File path (str) or PIL.Image object.

        Returns:
            dict: { "Class_Name": probability (0.0 - 1.0) }
        """
        try:
            # 1. Flexible Input Handling (Path or Object)
            if isinstance(image_source, str):
                image = Image.open(image_source).convert('RGB')
            else:
                image = image_source.convert('RGB')

            # 2. Preprocessing
            inputs = self.processor(images=image, return_tensors="pt")
            pixel_values = inputs['pixel_values'].to(self.device)

            # 3. Inference
            with torch.no_grad():
                # A. Get Embeddings from DINO
                outputs = self.base_model(pixel_values=pixel_values)
                
                # Handle different transformer output formats
                if hasattr(outputs, 'last_hidden_state'):
                    embeddings = outputs.last_hidden_state.mean(dim=1)
                else:
                    embeddings = outputs[0].mean(dim=1)

                # B. Classify Embeddings
                logits = self.classifier(embeddings)
                
                # Convert to standard Python float list for JSON serialization
                probs = torch.sigmoid(logits).cpu().numpy()[0].tolist()

            # 4. Format Output
            return {
                label: round(prob, 4) 
                for label, prob in zip(self.label_cols, probs)
            }

        except Exception as e:
            logger.error(f"Error predicting image: {e}")
            return {"error": str(e)}


# ==============================================================================
# 3. Factory Function (The "Builder")
# ==============================================================================
def create_classifier(checkpoint_path, model_id=_BUNDLED_VISION_MODEL, device=None):
    """
    Loads the checkpoint, reconstructs the specific architecture, 
    and returns a ready-to-use ChestXrayPredictor instance.

    Args:
        checkpoint_path (str): Path to the .pth file.
        model_id (str): HuggingFace model ID for DINO.
        device (str, optional): 'cuda' or 'cpu'. Auto-detects if None.
    
    Returns:
        ChestXrayPredictor: Initialized object ready for prediction.
    """
    device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info(f"🔄 Starting model initialization on: {device}")

    try:
        # A. Load Checkpoint
        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        label_cols = checkpoint.get('label_cols', [
            "Cardiomegaly", "Consolidation", "Edema",
            "Atelectasis", "Pleural Effusion", "No Findings"
        ])
        
        # B. Load Base Model (DINO)
        logger.info("🤖 Loading DINO backbone...")
        base_model = AutoModel.from_pretrained(model_id).to(device)
        
        # Load fine-tuned DINO weights if they exist in checkpoint
        if 'base_model_state_dict' in checkpoint:
            base_model.load_state_dict(checkpoint['base_model_state_dict'])
            logger.info("   - Fine-tuned DINO weights loaded from checkpoint.")
        else:
            logger.info("   - Using default pre-trained DINO weights.")
        
        processor = AutoImageProcessor.from_pretrained(model_id)

        # C. Detect Embedding Dimension
        if hasattr(base_model.config, 'hidden_size'):
            embedding_dim = base_model.config.hidden_size
        else:
            # Dummy inference to detect output size
            with torch.no_grad():
                dummy = torch.randn(1, 3, 224, 224).to(device)
                out = base_model(pixel_values=dummy)
                embedding_dim = out.last_hidden_state.shape[-1]

        # D. Reconstruct Classifier Architecture
        logger.info("🏗️  Reconstructing classifier architecture...")
        model_state = checkpoint['model_state_dict']
        classifier = _build_mlp_from_state(model_state, embedding_dim)
        
        # Load classifier weights
        classifier.load_state_dict(model_state)
        classifier.to(device)
        
        logger.info("✅ Model created successfully.")

        # E. Return the Wrapper Instance
        return ChestXrayPredictor(base_model, classifier, processor, label_cols, device)

    except Exception as e:
        logger.error(f"❌ Fatal error creating the classifier: {e}")
        raise e


def _build_mlp_from_state(model_state, embedding_dim):
    """
    Private helper function to inspect state_dict and rebuild the MLP architecture.
    """
    linear_layers = []
    for key, val in model_state.items():
        # Look for 2D weights (Linear layers) inside the classifier
        if 'classifier' in key and key.endswith('.weight') and len(val.shape) == 2:
            match = re.search(r'classifier\.(\d+)\.weight', key)
            if match:
                layer_idx = int(match.group(1))
                linear_layers.append((layer_idx, val.shape[1], val.shape[0])) # idx, in_features, out_features

    if not linear_layers:
        raise ValueError("No linear layers found in checkpoint. Check architecture.")
        
    # Sort by layer index to ensure correct order
    linear_layers.sort(key=lambda x: x[0])
    
    num_classes = linear_layers[-1][2]
    hidden_dims = tuple([x[2] for x in linear_layers[:-1]])
    
    # Detect Normalization types
    uses_bn = any('running_mean' in k for k in model_state.keys())
    has_norm = any(k.endswith('.weight') and len(model_state[k].shape) == 1 for k in model_state.keys() if 'classifier' in k)
    uses_layernorm = has_norm and not uses_bn

    return EmbeddingClassifier(
        embedding_dim=embedding_dim,
        num_classes=num_classes,
        custom_dims=hidden_dims,
        bn=uses_bn,
        use_layernorm=uses_layernorm
    )