Update utils/complete_model.py

utils/complete_model.py

@@ -1,276 +1,490 @@
-import os
-import torch
-import torch.nn as nn
-from transformers import AutoModel, GPT2Tokenizer
-
-from utils.modifiedGPT2 import create_decoder
-
-from utils.layer_mask import gaussian_layer_stack_pipeline
-
-
-class DINOEncoder(nn.Module):
-    def __init__(self, model_id="facebook/dinov3-vits16-pretrain-lvd1689m", freeze=True):
-        super().__init__()
-        self.model = AutoModel.from_pretrained(model_id)
-        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="facebook/dinov3-convnext-small-pretrain-lvd1689m", freeze=True):
-        super().__init__()
-        self.encoder = AutoModel.from_pretrained(model_name)
-        # 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("gpt2")
-        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
+import os
+import torch
+import torch.nn as nn
+from transformers import AutoModel, GPT2Tokenizer
+
+from utils.modifiedGPT2 import create_decoder
+
+from utils.layer_mask import gaussian_layer_stack_pipeline
+
+class DINOEncoder(nn.Module):
+    def __init__(self, model_id="facebook/dinov3-vits16-pretrain-lvd1689m", freeze=True):
+        super().__init__()
+        self.model = AutoModel.from_pretrained(model_id)
+        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="facebook/dinov3-convnext-small-pretrain-lvd1689m", freeze=True):
+        super().__init__()
+        self.encoder = AutoModel.from_pretrained(model_name)
+        # 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("gpt2")
+        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="facebook/dinov3-vits16-pretrain-lvd1689m", 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', ["Class_1", "Class_2"]) # Fallback
+        
+        # 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
+    )
+