lib/firefly_thermal.py

import os
import math
from functools import reduce
from operator import mul

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import SemanticSegmenterOutput

# ===================================================================
# 1. SHARED UTILITIES & LOSS FUNCTIONS
# ===================================================================

def dice_loss(pred, target, skip_classes=None, eps=1e-6):
    N, C = pred.shape
    
    if N == 0:
        return pred.sum() * 0.0

    p = torch.softmax(pred, dim=1)                        

    t = torch.nn.functional.one_hot(target, C).float()    

    if skip_classes is None:
        skip_classes = []

    present = (t.sum(dim=0) > 0)

    for sc in skip_classes:
        if 0 <= sc < C:
            present[sc] = False

    if present.sum() == 0:
        return pred.new_tensor(0.)

    p_sel = p[:, present]
    t_sel = t[:, present]

    inter = (p_sel * t_sel).sum(dim=0)                     
    union = p_sel.sum(dim=0) + t_sel.sum(dim=0)         
    dice = (2 * inter + eps) / (union + eps)              

    return 1 - dice.mean()

class DiceLoss(nn.Module):
    def __init__(self, skip_classes=None, eps=1e-6):
        super(DiceLoss, self).__init__()
        self.skip_classes = skip_classes
        self.eps = eps

    def forward(self, pred, target):
        return dice_loss(pred, target, skip_classes=self.skip_classes, eps=self.eps)

def pad_to_multiple(x, multiple):
    """Pads the input tensor so its spatial dimensions are multiples of 'multiple'."""
    _, _, h, w = x.shape
    pad_h = (multiple - h % multiple) % multiple
    pad_w = (multiple - w % multiple) % multiple
    padding = (0, pad_w, 0, pad_h)  
    x_padded = F.pad(x, padding, mode='reflect')
    return x_padded, padding

def crop_to_shape(x, target_h, target_w):
    """Crops the input tensor back to the target spatial dimensions."""
    return x[:, :, :target_h, :target_w]

class MLPHead(nn.Module):
    def __init__(self, in_channels=768, hidden_dim=256, num_classes=12, dropout_prob=0.1):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Dropout2d(p=dropout_prob),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Dropout2d(p=dropout_prob),
            nn.Conv2d(hidden_dim, num_classes, kernel_size=1)
        )
        self.upscale = True

    def forward(self, x):
        return self.model(x)

# ===================================================================
# 2. REIN IMPLEMENTATION
# ===================================================================

class Reins(nn.Module):
    def __init__(
        self,
        num_layers: int,
        embed_dims: int,
        patch_size: int,
        query_dims: int = 256,
        token_length: int = 100,
        use_softmax: bool = True,
        scale_init: float = 0.001,
    ) -> None:
        super().__init__()
        self.num_layers = num_layers
        self.embed_dims = embed_dims
        self.patch_size = patch_size
        self.query_dims = query_dims
        self.token_length = token_length
        self.scale_init = scale_init
        self.use_softmax = use_softmax
        self.create_model()

    def create_model(self):
        self.learnable_tokens = nn.Parameter(
            torch.empty([self.num_layers, self.token_length, self.embed_dims])
        )
        self.scale = nn.Parameter(torch.tensor(self.scale_init))
        
        self.mlp_token2feat = nn.Linear(self.embed_dims, self.embed_dims)
        self.mlp_delta_f = nn.Linear(self.embed_dims, self.embed_dims)
        
        val = math.sqrt(
            6.0
            / float(
                3 * reduce(mul, (self.patch_size, self.patch_size), 1) + self.embed_dims
            )
        )
        nn.init.uniform_(self.learnable_tokens.data, -val, val)
        nn.init.kaiming_uniform_(self.mlp_delta_f.weight, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.mlp_token2feat.weight, a=math.sqrt(5))

        self.transform = nn.Linear(self.embed_dims, self.query_dims)
        self.merge = nn.Linear(self.query_dims * 3, self.query_dims)

    def get_tokens(self, layer: int) -> Tensor:
        if layer == -1:
            return self.learnable_tokens
        else:
            return self.learnable_tokens[layer]

    def forward(
        self, feats: Tensor, layer: int, batch_first=False, has_cls_token=True
    ) -> Tensor:
        if batch_first:
            feats = feats.permute(1, 0, 2) 
        
        if has_cls_token:
            cls_token, feats = torch.tensor_split(feats, [1], dim=0)
            
        tokens = self.get_tokens(layer)
        
        delta_feat = self.forward_delta_feat(
            feats,
            tokens,
            layer,
        )
        
        delta_feat = delta_feat * self.scale
        feats = feats + delta_feat 
        
        if has_cls_token:
            feats = torch.cat([cls_token, feats], dim=0)
            
        if batch_first:
            feats = feats.permute(1, 0, 2)
            
        return feats

    def forward_delta_feat(self, feats: Tensor, tokens: Tensor, layers: int) -> Tensor:
        attn = torch.einsum("nbc,mc->nbm", feats, tokens)
        
        if self.use_softmax:
            attn = attn * (self.embed_dims**-0.5)
            attn = F.softmax(attn, dim=-1)
            
        delta_f = torch.einsum(
            "nbm,mc->nbc",
            attn[:, :, 1:], 
            self.mlp_token2feat(tokens[1:, :]),
        )
        
        delta_f = self.mlp_delta_f(delta_f + feats)
        return delta_f

class ReinHook:
    def __init__(self, rein_module, layer_idx):
        self.rein_module = rein_module
        self.layer_idx = layer_idx

    def __call__(self, module, input, output):
        is_tuple = isinstance(output, tuple)
        is_list = isinstance(output, list)
        
        if is_tuple:
            x = output[0]
        elif is_list:
            x = output[0]
        else:
            x = output

        refined_x = self.rein_module(x, self.layer_idx, batch_first=True, has_cls_token=True)
        
        if is_tuple:
            return (refined_x,) + output[1:]
        elif is_list:
            output[0] = refined_x
            return output
        else:
            return refined_x

class FireflyConfigThermal(PretrainedConfig):
    model_type = "firefly-thermal"

    def __init__(
        self,
        num_labels: int = 2,
        image_size: int = 512,
        embedding_dim: int = 256,
        backbone_embed_dim: int = 768,
        patch_size: int = 16,
        num_layers: int = 12,
        rein_token_length: int = 100,  
        feature_layers: list = None,
        repo_dir: str = None,
        model_name: str = "dinov3_vitb16",
        backbone_weights_path: str = None, 
        finetuned_weights_path: str = None, 
        semantic_loss_ignore_index: int = 255,
        dropout_ratio: float = 0.1,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.num_labels = num_labels
        self.image_size = image_size
        self.embedding_dim = embedding_dim
        self.backbone_embed_dim = backbone_embed_dim
        self.patch_size = patch_size
        self.num_layers = num_layers
        self.rein_token_length = rein_token_length
        self.feature_layers = feature_layers or [2, 5, 8, 11]
        self.repo_dir = repo_dir
        self.model_name = model_name
        self.backbone_weights_path = backbone_weights_path
        self.finetuned_weights_path = finetuned_weights_path
        self.semantic_loss_ignore_index = semantic_loss_ignore_index
        self.dropout_ratio = dropout_ratio

class FireflyForSemanticSegmentationThermal(PreTrainedModel):
    config_class = FireflyConfigThermal

    def __init__(self, config: FireflyConfigThermal):
        super().__init__(config)
        self.config = config
        self.num_labels = config.num_labels
        
        print(f"Loading backbone architecture: {config.model_name}")
        self.backbone = torch.hub.load(
            config.repo_dir,
            config.model_name,
            source="local",
            weights=config.backbone_weights_path,
            pretrained=False
        )
    
        print(f"Injecting Rein Module (Tokens={config.rein_token_length})...")
        self.rein = Reins(
            num_layers=config.num_layers,
            embed_dims=config.backbone_embed_dim,
            patch_size=config.patch_size,
            token_length=config.rein_token_length
        )

        if hasattr(self.backbone, 'blocks'):
            blocks = self.backbone.blocks
        elif hasattr(self.backbone, 'transformer') and hasattr(self.backbone.transformer, 'blocks'): 
             blocks = self.backbone.transformer.blocks
        else:
            raise AttributeError("Could not find '.blocks' in backbone model. Check model structure.")

        self.hooks = []
        for i, block in enumerate(blocks):
            hook_fn = ReinHook(self.rein, i)
            handle = block.register_forward_hook(hook_fn)
            self.hooks.append(handle)
        
        if config.finetuned_weights_path and os.path.exists(config.finetuned_weights_path):
            print(f"Loading finetuned rein adapter weights: {config.finetuned_weights_path}")
            checkpoint = torch.load(config.finetuned_weights_path, map_location='cpu')
            
            state_dict = checkpoint
            if isinstance(checkpoint, dict):
                if "rein_model" in checkpoint:
                    state_dict = checkpoint["rein_model"]
                elif "student_model" in checkpoint:
                    state_dict = checkpoint["student_model"]
                elif "state_dict" in checkpoint:
                    state_dict = checkpoint["state_dict"]
                elif "model" in checkpoint:
                    state_dict = checkpoint["model"]

            new_state_dict = {}
            for k, v in state_dict.items():
                clean_k = k
                for prefix in ["module.", "rein_model.", "rein.", "student_model.", "model."]:
                    if clean_k.startswith(prefix):
                        clean_k = clean_k.replace(prefix, "", 1)
                new_state_dict[clean_k] = v

            msg = self.rein.load_state_dict(new_state_dict, strict=False)
            print(msg)
            if len(msg.missing_keys) > 0:
                print(f"[Diagnostic] Found {len(msg.missing_keys)} missing keys in Rein module loading.")
                print(f"[Warning] Missing keys: {msg.missing_keys}")
            else:
                print("[Success] All Rein Adapter weights perfectly matched and loaded!")
        else:
            print(f"Warning: Finetuned weights not found at {config.finetuned_weights_path}. Using base/random weights.")

        self.head = MLPHead(
            in_channels=config.backbone_embed_dim,
            hidden_dim=config.embedding_dim,
            num_classes=config.num_labels,
            dropout_prob=config.dropout_ratio
        )
        
        self.loss_fn = DiceLoss()
        
        self.feature_layers = config.feature_layers
        self._setup_trainable_params()

    def _setup_trainable_params(self):
        for param in self.backbone.parameters():
            param.requires_grad = False
        
        for param in self.rein.parameters():
            param.requires_grad = True
        
        for param in self.head.parameters():
            param.requires_grad = True

    def _extract_features(self, pixel_values: torch.Tensor, h_pad: int, w_pad: int):
        features_raw = self.backbone.get_intermediate_layers(
            pixel_values, 
            n=1 
        )
        
        feat = features_raw[0]
        if isinstance(feat, tuple):
            feat = feat[0]

        if feat.ndim == 3:
            B, N, Dim = feat.shape
            H_grid = h_pad // self.config.patch_size
            W_grid = w_pad // self.config.patch_size
            
            num_spatial_tokens = H_grid * W_grid
            if N > num_spatial_tokens:
                feat = feat[:, -num_spatial_tokens:, :]
                
            feat = feat.permute(0, 2, 1).view(B, Dim, H_grid, W_grid)
            
        return feat

    def forward(self, pixel_values: torch.Tensor, labels: torch.Tensor = None, **kwargs):
        original_h, original_w = pixel_values.shape[-2:]
        
        x_padded, padding = pad_to_multiple(pixel_values, self.config.patch_size)
        pad_h, pad_w = x_padded.shape[-2], x_padded.shape[-1]
        
        features = self._extract_features(x_padded, pad_h, pad_w)
        out = self.head(features)
        
        out_upscaled = F.interpolate(
            out, 
            size=(pad_h, pad_w), 
            mode='bilinear', 
            align_corners=False
        )
        
        logits = crop_to_shape(out_upscaled, original_h, original_w)
        
        loss = None
        if labels is not None:
            valid_mask = labels != self.config.semantic_loss_ignore_index
            logits_masked = logits.permute(0, 2, 3, 1)[valid_mask]
            labels_masked = labels[valid_mask]
            loss = self.loss_fn(logits_masked, labels_masked)
        
        return SemanticSegmenterOutput(
            loss=loss,
            logits=logits,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path=None, **kwargs):
        config_kwargs = {
            "num_labels": kwargs.pop("num_labels", 2),
            "image_size": kwargs.pop("image_size", 512),
            "embedding_dim": kwargs.pop("embedding_dim", 256),
            "backbone_embed_dim": kwargs.pop("backbone_embed_dim", 768),
            "patch_size": kwargs.pop("patch_size", 16),
            "num_layers": kwargs.pop("num_layers", 12),
            "rein_token_length": kwargs.pop("rein_token_length", 100),
            "feature_layers": kwargs.pop("feature_layers", [2, 5, 8, 11]),
            "repo_dir": kwargs.pop("repo_dir", None),
            "model_name": kwargs.pop("model_name", "dinov3_vitb16"),
            "backbone_weights_path": kwargs.pop("backbone_weights_path", ""),
            "finetuned_weights_path": kwargs.pop("finetuned_weights_path", pretrained_model_name_or_path),
            "semantic_loss_ignore_index": kwargs.pop("semantic_loss_ignore_index", 255),
            "dropout_ratio": kwargs.pop("dropout_ratio", 0.1),
            "id2label": kwargs.pop("id2label", None),
            "label2id": kwargs.pop("label2id", None),
        }
        
        kwargs.pop("ignore_mismatched_sizes", None)
        config = kwargs.pop("config", None)
    
        if config is None:
            config = FireflyConfigThermal(**config_kwargs)
            
        model = cls(config)
        return model

    def print_trainable_params(self):
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        print(f"Total Parameters: {total / 1e6:.2f}M")
        print(f"Trainable Parameters: {trainable / 1e6:.2f}M")
        print(f"Trainable Ratio: {100 * trainable / total:.2f}%")