Update app/utils/model_loader.py

def predict function not being defined causes a syntax error when attempting to load the .safetensors model instead of the default onnx model

Files changed (1) hide show

app/utils/model_loader.py +379 -379

app/utils/model_loader.py CHANGED Viewed

@@ -1,379 +1,379 @@
-import torch
-import torch.nn as nn
-from torch.nn import GroupNorm, LayerNorm
-import torch.nn.functional as F
-import torch.utils.checkpoint as checkpoint
-import timm
-class ViTWrapper(nn.Module):
-    """Wrapper to make ViT compatible with feature extraction for ImageTagger"""
-    def __init__(self, vit_model):
-        super().__init__()
-        self.vit = vit_model
-        self.out_indices = (-1,)  # mimic timm.features_only
-        # Get patch size and embedding dim from the model
-        self.patch_size = vit_model.patch_embed.patch_size[0]
-        self.embed_dim = vit_model.embed_dim
-    def forward(self, x):
-        B = x.size(0)
-        # ➊ patch tokens
-        x = self.vit.patch_embed(x)                       # (B, N, C)
-        # ➋ prepend CLS
-        cls_tok = self.vit.cls_token.expand(B, -1, -1)    # (B, 1, C)
-        x = torch.cat((cls_tok, x), dim=1)                # (B, 1+N, C)
-        # ➌ add positional encodings (full, incl. CLS)
-        if self.vit.pos_embed is not None:
-            x = x + self.vit.pos_embed[:, : x.size(1), :]
-        x = self.vit.pos_drop(x)
-        for blk in self.vit.blocks:
-            x = blk(x)
-        x = self.vit.norm(x)                              # (B, 1+N, C)
-        # ➍ split back out
-        cls_final   = x[:, 0]              # (B, C)
-        patch_tokens = x[:, 1:]            # (B, N, C)
-        # ➎ reshape patches to (B, C, H, W)
-        B, N, C = patch_tokens.shape
-        h = w = int(N ** 0.5)              # square assumption
-        patch_features = patch_tokens.permute(0, 2, 1).reshape(B, C, h, w)
-        # Return **both**: (patch map, CLS)
-        return patch_features, cls_final
-    def set_grad_checkpointing(self, enable=True):
-        """Enable gradient checkpointing if supported"""
-        if hasattr(self.vit, 'set_grad_checkpointing'):
-            self.vit.set_grad_checkpointing(enable)
-            return True
-        return False
-class ImageTagger(nn.Module):
-    """
-    ImageTagger with Vision Transformer backbone
-    """
-    def __init__(self, total_tags, dataset, model_name='vit_base_patch16_224',
-                 num_heads=16, dropout=0.1, pretrained=True, tag_context_size=256,
-                 use_gradient_checkpointing=False, img_size=224):
-        super().__init__()
-        # Store checkpointing config
-        self.use_gradient_checkpointing = use_gradient_checkpointing
-        self.model_name = model_name
-        self.img_size = img_size
-        # Debug and stats flags
-        self._flags = {
-            'debug': False,
-            'model_stats': True
-        }
-        # Core model config
-        self.dataset = dataset
-        self.tag_context_size = tag_context_size
-        self.total_tags = total_tags
-        print(f"🏗️ Building ImageTagger with ViT backbone and {total_tags} tags")
-        print(f"   Backbone: {model_name}")
-        print(f"   Image size: {img_size}x{img_size}")
-        print(f"   Tag context size: {tag_context_size}")
-        print(f"   Gradient checkpointing: {use_gradient_checkpointing}")
-        print(f"   🎯 Custom embeddings, PyTorch native attention, no ground truth inclusion")
-        # 1. Vision Transformer Backbone
-        print("📦 Loading Vision Transformer backbone...")
-        self._load_vit_backbone()
-        # Get backbone dimensions by running a test forward pass
-        self._determine_backbone_dimensions()
-        self.embedding_dim = self.backbone.embed_dim
-        # 2. Custom Tag Embeddings (no CLIP)
-        print("🎯 Using custom tag embeddings (no CLIP)")
-        self.tag_embedding = nn.Embedding(total_tags, self.embedding_dim)
-        # 3. Shared weights approach - tag bias for initial predictions
-        print("🔗 Using shared weights between initial head and tag embeddings")
-        self.tag_bias = nn.Parameter(torch.zeros(total_tags))
-        # 4. Image token extraction (for attention AND global pooling)
-        self.image_token_proj = nn.Identity()
-        # 5. Tags-as-queries cross-attention (using PyTorch's optimized implementation)
-        self.cross_attention = nn.MultiheadAttention(
-            embed_dim=self.embedding_dim,
-            num_heads=num_heads,
-            dropout=dropout,
-            batch_first=True  # Use (batch, seq, feature) format
-        )
-        self.cross_norm = nn.LayerNorm(self.embedding_dim)
-        # Initialize weights
-        self._init_weights()
-        # Enable gradient checkpointing
-        if self.use_gradient_checkpointing:
-            self._enable_gradient_checkpointing()
-        print(f"✅ ImageTagger with ViT initialized!")
-        self._print_parameter_count()
-    def _load_vit_backbone(self):
-        """Load Vision Transformer model from timm"""
-        print(f"   Loading from timm: {self.model_name}")
-        # Load the ViT model (not features_only, we want the full model for token extraction)
-        vit_model = timm.create_model(
-            self.model_name,
-            pretrained=True,
-            img_size=self.img_size,
-            num_classes=0  # Remove classification head
-        )
-        # Wrap it in our compatibility layer
-        self.backbone = ViTWrapper(vit_model)
-        print(f"   ✅ ViT loaded successfully")
-        print(f"   Patch size: {self.backbone.patch_size}x{self.backbone.patch_size}")
-        print(f"   Embed dim: {self.backbone.embed_dim}")
-    def _determine_backbone_dimensions(self):
-        """Determine backbone output dimensions"""
-        print("   🔍 Determining backbone dimensions...")
-        with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
-            # Create a dummy input
-            dummy_input = torch.randn(1, 3, self.img_size, self.img_size)
-            # Get features
-            backbone_features, cls_dummy = self.backbone(dummy_input)
-            feature_tensor = backbone_features
-            self.backbone_dim = feature_tensor.shape[1]
-            self.feature_map_size = feature_tensor.shape[2]
-        print(f"   Backbone output: {self.backbone_dim}D, {self.feature_map_size}x{self.feature_map_size} spatial")
-        print(f"   Total patch tokens: {self.feature_map_size * self.feature_map_size}")
-    def _enable_gradient_checkpointing(self):
-        """Enable gradient checkpointing for memory efficiency"""
-        print("🔄 Enabling gradient checkpointing...")
-        # Enable checkpointing for ViT backbone
-        if self.backbone.set_grad_checkpointing(True):
-            print("   ✅ ViT backbone checkpointing enabled")
-        else:
-            print("   ⚠️ ViT backbone doesn't support built-in checkpointing, will checkpoint manually")
-    def _checkpoint_backbone(self, x):
-        """Wrapper for backbone with gradient checkpointing"""
-        if self.use_gradient_checkpointing and self.training:
-            return checkpoint.checkpoint(self.backbone, x, use_reentrant=False)
-        else:
-            return self.backbone(x)
-    def _checkpoint_image_proj(self, x):
-        """Wrapper for image projection with gradient checkpointing"""
-        if self.use_gradient_checkpointing and self.training:
-            return checkpoint.checkpoint(self.image_token_proj, x, use_reentrant=False)
-        else:
-            return self.image_token_proj(x)
-    def _checkpoint_cross_attention(self, query, key, value):
-        """Wrapper for cross attention with gradient checkpointing"""
-        def _attention_forward(q, k, v):
-            attended_features, _ = self.cross_attention(query=q, key=k, value=v)
-            return self.cross_norm(attended_features)
-        if self.use_gradient_checkpointing and self.training:
-            return checkpoint.checkpoint(_attention_forward, query, key, value, use_reentrant=False)
-        else:
-            return _attention_forward(query, key, value)
-    def _checkpoint_candidate_selection(self, initial_logits):
-        """Wrapper for candidate selection with gradient checkpointing"""
-        def _candidate_forward(logits):
-            return self._get_candidate_tags(logits)
-        if self.use_gradient_checkpointing and self.training:
-            return checkpoint.checkpoint(_candidate_forward, initial_logits, use_reentrant=False)
-        else:
-            return _candidate_forward(initial_logits)
-    def _checkpoint_final_scoring(self, attended_features, candidate_indices):
-        """Wrapper for final scoring with gradient checkpointing"""
-        def _scoring_forward(features, indices):
-            emb = self.tag_embedding(indices)
-            # BF16 in, BF16 out
-            return (features * emb).sum(dim=-1)
-        if self.use_gradient_checkpointing and self.training:
-            return checkpoint.checkpoint(_scoring_forward, attended_features, candidate_indices, use_reentrant=False)
-        else:
-            return _scoring_forward(attended_features, candidate_indices)
-    def _init_weights(self):
-        """Initialize weights for new modules"""
-        def _init_layer(layer):
-            if isinstance(layer, nn.Linear):
-                nn.init.xavier_uniform_(layer.weight)
-                if layer.bias is not None:
-                    nn.init.zeros_(layer.bias)
-            elif isinstance(layer, nn.Conv2d):
-                nn.init.kaiming_normal_(layer.weight, mode='fan_out', nonlinearity='relu')
-                if layer.bias is not None:
-                    nn.init.zeros_(layer.bias)
-            elif isinstance(layer, nn.Embedding):
-                nn.init.normal_(layer.weight, mean=0, std=0.02)
-        # Initialize new components
-        self.image_token_proj.apply(_init_layer)
-        # Initialize tag embeddings with normal distribution
-        nn.init.normal_(self.tag_embedding.weight, mean=0, std=0.02)
-        # Initialize tag bias
-        nn.init.zeros_(self.tag_bias)
-    def _print_parameter_count(self):
-        """Print parameter statistics"""
-        total_params = sum(p.numel() for p in self.parameters())
-        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
-        backbone_params = sum(p.numel() for p in self.backbone.parameters())
-        print(f"📊 Parameter Statistics:")
-        print(f"   Total parameters: {total_params/1e6:.1f}M")
-        print(f"   Trainable parameters: {trainable_params/1e6:.1f}M")
-        print(f"   Frozen parameters: {(total_params-trainable_params)/1e6:.1f}M")
-        print(f"   Backbone parameters: {backbone_params/1e6:.1f}M")
-        if self.use_gradient_checkpointing:
-            print(f"   🔄 Gradient checkpointing enabled for memory efficiency")
-    @property
-    def debug(self):
-        return self._flags['debug']
-    @property
-    def model_stats(self):
-        return self._flags['model_stats']
-    def _get_candidate_tags(self, initial_logits, target_tags=None, hard_negatives=None):
-        """Select candidate tags - no ground truth inclusion"""
-        batch_size = initial_logits.size(0)
-        # Simply select top K candidates based on initial predictions
-        top_probs, top_indices = torch.topk(
-            torch.sigmoid(initial_logits),
-            k=min(self.tag_context_size, self.total_tags),
-            dim=1, largest=True, sorted=True
-        )
-        return top_indices
-    def _analyze_predictions(self, predictions, tag_indices):
-        """Analyze prediction patterns"""
-        if not self.model_stats:
-            return {}
-        if torch._dynamo.is_compiling():
-            return {}
-        with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
-            probs = torch.sigmoid(predictions)
-            relevant_probs = torch.gather(probs, 1, tag_indices)
-            return {
-                'prediction_confidence': relevant_probs.mean().item(),
-                'prediction_entropy': -(relevant_probs * torch.log(relevant_probs + 1e-9)).mean().item(),
-                'high_confidence_ratio': (relevant_probs > 0.7).float().mean().item(),
-                'above_threshold_ratio': (relevant_probs > 0.5).float().mean().item(),
-            }
-    def forward(self, x, targets=None, hard_negatives=None):
-        """
-        Forward pass with ViT backbone, CLS token support and gradient-checkpointing.
-        All arithmetic tensors stay in the backbone’s dtype (BF16 under autocast,
-        FP32 otherwise).  Anything that must mix dtypes is cast to match.
-        """
-        batch_size  = x.size(0)
-        model_stats = {} if self.model_stats else {}
-        # ------------------------------------------------------------------
-        # 1. Backbone  →  patch map + CLS token
-        # ------------------------------------------------------------------
-        patch_map, cls_token = self._checkpoint_backbone(x)         # patch_map: [B, C, H, W]
-                                                                    # cls_token: [B, C]
-        # ------------------------------------------------------------------
-        # 2. Tokens  →  global image vector
-        # ------------------------------------------------------------------
-        image_tokens_4d = self._checkpoint_image_proj(patch_map)    # [B, C, H, W]
-        image_tokens    = image_tokens_4d.flatten(2).transpose(1, 2)  # [B, N, C]
-        # “Dual-pool”: mean-pool patches ⊕ CLS
-        global_features = 0.5 * (image_tokens.mean(dim=1, dtype=image_tokens.dtype) + cls_token)  # [B, C]
-        compute_dtype = global_features.dtype                       # BF16 or FP32
-        # ------------------------------------------------------------------
-        # 3. Initial logits  (shared weights)
-        # ------------------------------------------------------------------
-        tag_weights = self.tag_embedding.weight.to(compute_dtype)   # [T, C]
-        tag_bias    = self.tag_bias.to(compute_dtype)               # [T]
-        initial_logits = global_features @ tag_weights.t() + tag_bias   # [B, T]
-        initial_logits = initial_logits.to(compute_dtype)               # keep dtype uniform
-        initial_preds  = initial_logits                                 # alias
-        # ------------------------------------------------------------------
-        # 4. Candidate set
-        # ------------------------------------------------------------------
-        candidate_indices = self._checkpoint_candidate_selection(initial_logits)  # [B, K]
-        tag_embeddings   = self.tag_embedding(candidate_indices).to(compute_dtype)  # [B, K, C]
-        attended_features = self._checkpoint_cross_attention(       # [B, K, C]
-            tag_embeddings, image_tokens, image_tokens
-        )
-        # ------------------------------------------------------------------
-        # 5. Score candidates  &  scatter back
-        # ------------------------------------------------------------------
-        candidate_logits = self._checkpoint_final_scoring(attended_features, candidate_indices)  # [B, K]
-        # --- align dtypes so scatter never throws ---
-        if candidate_logits.dtype != initial_logits.dtype:
-            candidate_logits = candidate_logits.to(initial_logits.dtype)
-        refined_logits = initial_logits.clone()
-        refined_logits.scatter_(1, candidate_indices, candidate_logits)
-        refined_preds = refined_logits
-        # ------------------------------------------------------------------
-        # 6. Optional stats
-        # ------------------------------------------------------------------
-        if self.model_stats and targets is not None and not torch._dynamo.is_compiling():
-            model_stats['initial_prediction_stats'] = self._analyze_predictions(initial_preds,
-                                                                                candidate_indices)
-            model_stats['refined_prediction_stats'] = self._analyze_predictions(refined_preds,
-                                                                                candidate_indices)
-        return {
-            'initial_predictions': initial_preds,
-            'refined_predictions': refined_preds,
-            'selected_candidates': candidate_indices,
-            'model_stats': model_stats
-        }
-    def predict

+import torch
+import torch.nn as nn
+from torch.nn import GroupNorm, LayerNorm
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+import timm
+class ViTWrapper(nn.Module):
+    """Wrapper to make ViT compatible with feature extraction for ImageTagger"""
+    def __init__(self, vit_model):
+        super().__init__()
+        self.vit = vit_model
+        self.out_indices = (-1,)  # mimic timm.features_only
+        # Get patch size and embedding dim from the model
+        self.patch_size = vit_model.patch_embed.patch_size[0]
+        self.embed_dim = vit_model.embed_dim
+    def forward(self, x):
+        B = x.size(0)
+        # ➊ patch tokens
+        x = self.vit.patch_embed(x)                       # (B, N, C)
+        # ➋ prepend CLS
+        cls_tok = self.vit.cls_token.expand(B, -1, -1)    # (B, 1, C)
+        x = torch.cat((cls_tok, x), dim=1)                # (B, 1+N, C)
+        # ➌ add positional encodings (full, incl. CLS)
+        if self.vit.pos_embed is not None:
+            x = x + self.vit.pos_embed[:, : x.size(1), :]
+        x = self.vit.pos_drop(x)
+        for blk in self.vit.blocks:
+            x = blk(x)
+        x = self.vit.norm(x)                              # (B, 1+N, C)
+        # ➍ split back out
+        cls_final   = x[:, 0]              # (B, C)
+        patch_tokens = x[:, 1:]            # (B, N, C)
+        # ➎ reshape patches to (B, C, H, W)
+        B, N, C = patch_tokens.shape
+        h = w = int(N ** 0.5)              # square assumption
+        patch_features = patch_tokens.permute(0, 2, 1).reshape(B, C, h, w)
+        # Return **both**: (patch map, CLS)
+        return patch_features, cls_final
+    def set_grad_checkpointing(self, enable=True):
+        """Enable gradient checkpointing if supported"""
+        if hasattr(self.vit, 'set_grad_checkpointing'):
+            self.vit.set_grad_checkpointing(enable)
+            return True
+        return False
+class ImageTagger(nn.Module):
+    """
+    ImageTagger with Vision Transformer backbone
+    """
+    def __init__(self, total_tags, dataset, model_name='vit_base_patch16_224',
+                 num_heads=16, dropout=0.1, pretrained=True, tag_context_size=256,
+                 use_gradient_checkpointing=False, img_size=224):
+        super().__init__()
+        # Store checkpointing config
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.model_name = model_name
+        self.img_size = img_size
+        # Debug and stats flags
+        self._flags = {
+            'debug': False,
+            'model_stats': True
+        }
+        # Core model config
+        self.dataset = dataset
+        self.tag_context_size = tag_context_size
+        self.total_tags = total_tags
+        print(f"🏗️ Building ImageTagger with ViT backbone and {total_tags} tags")
+        print(f"   Backbone: {model_name}")
+        print(f"   Image size: {img_size}x{img_size}")
+        print(f"   Tag context size: {tag_context_size}")
+        print(f"   Gradient checkpointing: {use_gradient_checkpointing}")
+        print(f"   🎯 Custom embeddings, PyTorch native attention, no ground truth inclusion")
+        # 1. Vision Transformer Backbone
+        print("📦 Loading Vision Transformer backbone...")
+        self._load_vit_backbone()
+        # Get backbone dimensions by running a test forward pass
+        self._determine_backbone_dimensions()
+        self.embedding_dim = self.backbone.embed_dim
+        # 2. Custom Tag Embeddings (no CLIP)
+        print("🎯 Using custom tag embeddings (no CLIP)")
+        self.tag_embedding = nn.Embedding(total_tags, self.embedding_dim)
+        # 3. Shared weights approach - tag bias for initial predictions
+        print("🔗 Using shared weights between initial head and tag embeddings")
+        self.tag_bias = nn.Parameter(torch.zeros(total_tags))
+        # 4. Image token extraction (for attention AND global pooling)
+        self.image_token_proj = nn.Identity()
+        # 5. Tags-as-queries cross-attention (using PyTorch's optimized implementation)
+        self.cross_attention = nn.MultiheadAttention(
+            embed_dim=self.embedding_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True  # Use (batch, seq, feature) format
+        )
+        self.cross_norm = nn.LayerNorm(self.embedding_dim)
+        # Initialize weights
+        self._init_weights()
+        # Enable gradient checkpointing
+        if self.use_gradient_checkpointing:
+            self._enable_gradient_checkpointing()
+        print(f"✅ ImageTagger with ViT initialized!")
+        self._print_parameter_count()
+    def _load_vit_backbone(self):
+        """Load Vision Transformer model from timm"""
+        print(f"   Loading from timm: {self.model_name}")
+        # Load the ViT model (not features_only, we want the full model for token extraction)
+        vit_model = timm.create_model(
+            self.model_name,
+            pretrained=True,
+            img_size=self.img_size,
+            num_classes=0  # Remove classification head
+        )
+        # Wrap it in our compatibility layer
+        self.backbone = ViTWrapper(vit_model)
+        print(f"   ✅ ViT loaded successfully")
+        print(f"   Patch size: {self.backbone.patch_size}x{self.backbone.patch_size}")
+        print(f"   Embed dim: {self.backbone.embed_dim}")
+    def _determine_backbone_dimensions(self):
+        """Determine backbone output dimensions"""
+        print("   🔍 Determining backbone dimensions...")
+        with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
+            # Create a dummy input
+            dummy_input = torch.randn(1, 3, self.img_size, self.img_size)
+            # Get features
+            backbone_features, cls_dummy = self.backbone(dummy_input)
+            feature_tensor = backbone_features
+            self.backbone_dim = feature_tensor.shape[1]
+            self.feature_map_size = feature_tensor.shape[2]
+        print(f"   Backbone output: {self.backbone_dim}D, {self.feature_map_size}x{self.feature_map_size} spatial")
+        print(f"   Total patch tokens: {self.feature_map_size * self.feature_map_size}")
+    def _enable_gradient_checkpointing(self):
+        """Enable gradient checkpointing for memory efficiency"""
+        print("🔄 Enabling gradient checkpointing...")
+        # Enable checkpointing for ViT backbone
+        if self.backbone.set_grad_checkpointing(True):
+            print("   ✅ ViT backbone checkpointing enabled")
+        else:
+            print("   ⚠️ ViT backbone doesn't support built-in checkpointing, will checkpoint manually")
+    def _checkpoint_backbone(self, x):
+        """Wrapper for backbone with gradient checkpointing"""
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(self.backbone, x, use_reentrant=False)
+        else:
+            return self.backbone(x)
+    def _checkpoint_image_proj(self, x):
+        """Wrapper for image projection with gradient checkpointing"""
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(self.image_token_proj, x, use_reentrant=False)
+        else:
+            return self.image_token_proj(x)
+    def _checkpoint_cross_attention(self, query, key, value):
+        """Wrapper for cross attention with gradient checkpointing"""
+        def _attention_forward(q, k, v):
+            attended_features, _ = self.cross_attention(query=q, key=k, value=v)
+            return self.cross_norm(attended_features)
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(_attention_forward, query, key, value, use_reentrant=False)
+        else:
+            return _attention_forward(query, key, value)
+    def _checkpoint_candidate_selection(self, initial_logits):
+        """Wrapper for candidate selection with gradient checkpointing"""
+        def _candidate_forward(logits):
+            return self._get_candidate_tags(logits)
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(_candidate_forward, initial_logits, use_reentrant=False)
+        else:
+            return _candidate_forward(initial_logits)
+    def _checkpoint_final_scoring(self, attended_features, candidate_indices):
+        """Wrapper for final scoring with gradient checkpointing"""
+        def _scoring_forward(features, indices):
+            emb = self.tag_embedding(indices)
+            # BF16 in, BF16 out
+            return (features * emb).sum(dim=-1)
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(_scoring_forward, attended_features, candidate_indices, use_reentrant=False)
+        else:
+            return _scoring_forward(attended_features, candidate_indices)
+    def _init_weights(self):
+        """Initialize weights for new modules"""
+        def _init_layer(layer):
+            if isinstance(layer, nn.Linear):
+                nn.init.xavier_uniform_(layer.weight)
+                if layer.bias is not None:
+                    nn.init.zeros_(layer.bias)
+            elif isinstance(layer, nn.Conv2d):
+                nn.init.kaiming_normal_(layer.weight, mode='fan_out', nonlinearity='relu')
+                if layer.bias is not None:
+                    nn.init.zeros_(layer.bias)
+            elif isinstance(layer, nn.Embedding):
+                nn.init.normal_(layer.weight, mean=0, std=0.02)
+        # Initialize new components
+        self.image_token_proj.apply(_init_layer)
+        # Initialize tag embeddings with normal distribution
+        nn.init.normal_(self.tag_embedding.weight, mean=0, std=0.02)
+        # Initialize tag bias
+        nn.init.zeros_(self.tag_bias)
+    def _print_parameter_count(self):
+        """Print parameter statistics"""
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        backbone_params = sum(p.numel() for p in self.backbone.parameters())
+        print(f"📊 Parameter Statistics:")
+        print(f"   Total parameters: {total_params/1e6:.1f}M")
+        print(f"   Trainable parameters: {trainable_params/1e6:.1f}M")
+        print(f"   Frozen parameters: {(total_params-trainable_params)/1e6:.1f}M")
+        print(f"   Backbone parameters: {backbone_params/1e6:.1f}M")
+        if self.use_gradient_checkpointing:
+            print(f"   🔄 Gradient checkpointing enabled for memory efficiency")
+    @property
+    def debug(self):
+        return self._flags['debug']
+    @property
+    def model_stats(self):
+        return self._flags['model_stats']
+    def _get_candidate_tags(self, initial_logits, target_tags=None, hard_negatives=None):
+        """Select candidate tags - no ground truth inclusion"""
+        batch_size = initial_logits.size(0)
+        # Simply select top K candidates based on initial predictions
+        top_probs, top_indices = torch.topk(
+            torch.sigmoid(initial_logits),
+            k=min(self.tag_context_size, self.total_tags),
+            dim=1, largest=True, sorted=True
+        )
+        return top_indices
+    def _analyze_predictions(self, predictions, tag_indices):
+        """Analyze prediction patterns"""
+        if not self.model_stats:
+            return {}
+        if torch._dynamo.is_compiling():
+            return {}
+        with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
+            probs = torch.sigmoid(predictions)
+            relevant_probs = torch.gather(probs, 1, tag_indices)
+            return {
+                'prediction_confidence': relevant_probs.mean().item(),
+                'prediction_entropy': -(relevant_probs * torch.log(relevant_probs + 1e-9)).mean().item(),
+                'high_confidence_ratio': (relevant_probs > 0.7).float().mean().item(),
+                'above_threshold_ratio': (relevant_probs > 0.5).float().mean().item(),
+            }
+    def forward(self, x, targets=None, hard_negatives=None):
+        """
+        Forward pass with ViT backbone, CLS token support and gradient-checkpointing.
+        All arithmetic tensors stay in the backbone’s dtype (BF16 under autocast,
+        FP32 otherwise).  Anything that must mix dtypes is cast to match.
+        """
+        batch_size  = x.size(0)
+        model_stats = {} if self.model_stats else {}
+        # ------------------------------------------------------------------
+        # 1. Backbone  →  patch map + CLS token
+        # ------------------------------------------------------------------
+        patch_map, cls_token = self._checkpoint_backbone(x)         # patch_map: [B, C, H, W]
+                                                                    # cls_token: [B, C]
+        # ------------------------------------------------------------------
+        # 2. Tokens  →  global image vector
+        # ------------------------------------------------------------------
+        image_tokens_4d = self._checkpoint_image_proj(patch_map)    # [B, C, H, W]
+        image_tokens    = image_tokens_4d.flatten(2).transpose(1, 2)  # [B, N, C]
+        # “Dual-pool”: mean-pool patches ⊕ CLS
+        global_features = 0.5 * (image_tokens.mean(dim=1, dtype=image_tokens.dtype) + cls_token)  # [B, C]
+        compute_dtype = global_features.dtype                       # BF16 or FP32
+        # ------------------------------------------------------------------
+        # 3. Initial logits  (shared weights)
+        # ------------------------------------------------------------------
+        tag_weights = self.tag_embedding.weight.to(compute_dtype)   # [T, C]
+        tag_bias    = self.tag_bias.to(compute_dtype)               # [T]
+        initial_logits = global_features @ tag_weights.t() + tag_bias   # [B, T]
+        initial_logits = initial_logits.to(compute_dtype)               # keep dtype uniform
+        initial_preds  = initial_logits                                 # alias
+        # ------------------------------------------------------------------
+        # 4. Candidate set
+        # ------------------------------------------------------------------
+        candidate_indices = self._checkpoint_candidate_selection(initial_logits)  # [B, K]
+        tag_embeddings   = self.tag_embedding(candidate_indices).to(compute_dtype)  # [B, K, C]
+        attended_features = self._checkpoint_cross_attention(       # [B, K, C]
+            tag_embeddings, image_tokens, image_tokens
+        )
+        # ------------------------------------------------------------------
+        # 5. Score candidates  &  scatter back
+        # ------------------------------------------------------------------
+        candidate_logits = self._checkpoint_final_scoring(attended_features, candidate_indices)  # [B, K]
+        # --- align dtypes so scatter never throws ---
+        if candidate_logits.dtype != initial_logits.dtype:
+            candidate_logits = candidate_logits.to(initial_logits.dtype)
+        refined_logits = initial_logits.clone()
+        refined_logits.scatter_(1, candidate_indices, candidate_logits)
+        refined_preds = refined_logits
+        # ------------------------------------------------------------------
+        # 6. Optional stats
+        # ------------------------------------------------------------------
+        if self.model_stats and targets is not None and not torch._dynamo.is_compiling():
+            model_stats['initial_prediction_stats'] = self._analyze_predictions(initial_preds,
+                                                                                candidate_indices)
+            model_stats['refined_prediction_stats'] = self._analyze_predictions(refined_preds,
+                                                                                candidate_indices)
+        return {
+            'initial_predictions': initial_preds,
+            'refined_predictions': refined_preds,
+            'selected_candidates': candidate_indices,
+            'model_stats': model_stats
+        }
+   # def predict