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}%")