model_code.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import efficientnet_v2_l, EfficientNet_V2_L_Weights
from PIL import Image
from typing import Optional
import torchvision.transforms as transforms
import os
import json

class InitialOnlyImageTagger(nn.Module):
    """
    A lightweight version of ImageTagger that only includes the backbone and initial classifier.
    This model uses significantly less VRAM than the full model.
    """
    def __init__(self, total_tags, dataset, model_name='efficientnet_v2_l',
                 dropout=0.1, pretrained=True):
        super().__init__()
        # Debug and stats flags
        self._flags = {
            'debug': False,
            'model_stats': False
        }
        
        # Core model config
        self.dataset = dataset
        self.embedding_dim = 1280  # Fixed to EfficientNetV2-L output dimension
        
        # Initialize backbone
        if model_name == 'efficientnet_v2_l':
            weights = EfficientNet_V2_L_Weights.DEFAULT if pretrained else None
            self.backbone = efficientnet_v2_l(weights=weights)
            self.backbone.classifier = nn.Identity()
        
        # Spatial pooling only - no projection
        self.spatial_pool = nn.AdaptiveAvgPool2d((1, 1))
        
        # Initial tag prediction with bottleneck
        self.initial_classifier = nn.Sequential(
            nn.Linear(self.embedding_dim, self.embedding_dim * 2),
            nn.LayerNorm(self.embedding_dim * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(self.embedding_dim * 2, self.embedding_dim),
            nn.LayerNorm(self.embedding_dim),
            nn.GELU(),
            nn.Linear(self.embedding_dim, total_tags)
        )
        
        # Temperature scaling
        self.temperature = nn.Parameter(torch.ones(1) * 1.5)
        
    @property
    def debug(self):
        return self._flags['debug']
    
    @debug.setter
    def debug(self, value):
        self._flags['debug'] = value
        
    @property
    def model_stats(self):
        return self._flags['model_stats']
    
    @model_stats.setter
    def model_stats(self, value):
        self._flags['model_stats'] = value
        
    def preprocess_image(self, image_path, image_size=512):
        """Process an image for inference using same preprocessing as training"""
        if not os.path.exists(image_path):
            raise ValueError(f"Image not found at path: {image_path}")
       
        # Initialize the same transform used during training
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
       
        try:
            with Image.open(image_path) as img:
                # Convert RGBA or Palette images to RGB
                if img.mode in ('RGBA', 'P'):
                    img = img.convert('RGB')
               
                # Get original dimensions
                width, height = img.size
                aspect_ratio = width / height
               
                # Calculate new dimensions to maintain aspect ratio
                if aspect_ratio > 1:
                    new_width = image_size
                    new_height = int(new_width / aspect_ratio)
                else:
                    new_height = image_size
                    new_width = int(new_height * aspect_ratio)
               
                # Resize with LANCZOS filter
                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
               
                # Create new image with padding
                new_image = Image.new('RGB', (image_size, image_size), (0, 0, 0))
                paste_x = (image_size - new_width) // 2
                paste_y = (image_size - new_height) // 2
                new_image.paste(img, (paste_x, paste_y))
               
                # Apply transforms (without normalization)
                img_tensor = transform(new_image)
                return img_tensor
        except Exception as e:
            raise Exception(f"Error processing {image_path}: {str(e)}")
    
    def forward(self, x):
        """Forward pass with only the initial predictions"""
        # Image Feature Extraction
        features = self.backbone.features(x)
        features = self.spatial_pool(features).squeeze(-1).squeeze(-1)
        
        # Initial Tag Predictions
        initial_logits = self.initial_classifier(features)
        initial_preds = torch.clamp(initial_logits / self.temperature, min=-15.0, max=15.0)
        
        # For API compatibility with the full model, return the same predictions twice
        return initial_preds, initial_preds
    
    def predict(self, image_path, threshold=0.325, category_thresholds=None):
        """
        Run inference on an image with support for category-specific thresholds.
        """
        # Preprocess the image
        img_tensor = self.preprocess_image(image_path).unsqueeze(0)
        
        # Move to the same device as model and convert to half precision
        device = next(self.parameters()).device
        dtype = next(self.parameters()).dtype  # Match model's precision
        img_tensor = img_tensor.to(device, dtype=dtype)
        
        # Run inference
        with torch.no_grad():
            initial_preds, _ = self.forward(img_tensor)
            
            # Apply sigmoid to get probabilities
            initial_probs = torch.sigmoid(initial_preds)
            
            # Apply thresholds
            if category_thresholds:
                # Create binary prediction tensors
                initial_binary = torch.zeros_like(initial_probs)
                
                # Apply thresholds by category
                for category, cat_threshold in category_thresholds.items():
                    # Create a mask for tags in this category
                    category_mask = torch.zeros_like(initial_probs, dtype=torch.bool)
                    
                    # Find indices for this category
                    for tag_idx in range(initial_probs.size(-1)):
                        try:
                            _, tag_category = self.dataset.get_tag_info(tag_idx)
                            if tag_category == category:
                                category_mask[:, tag_idx] = True
                        except:
                            continue
                    
                    # Apply threshold only to tags in this category
                    cat_threshold_tensor = torch.tensor(cat_threshold, device=device, dtype=dtype)
                    initial_binary[category_mask] = (initial_probs[category_mask] >= cat_threshold_tensor).to(dtype)
                
                predictions = initial_binary
            else:
                # Use the same threshold for all tags
                threshold_tensor = torch.tensor(threshold, device=device, dtype=dtype)
                predictions = (initial_probs >= threshold_tensor).to(dtype)
            
            # Return the same probabilities for both initial and refined for API compatibility
            return {
                'initial_probabilities': initial_probs,
                'refined_probabilities': initial_probs,  # Same as initial for compatibility
                'predictions': predictions
            }
            
    def get_tags_from_predictions(self, predictions, include_probabilities=True):
        """
        Convert model predictions to human-readable tags grouped by category.
        """
        # Get non-zero predictions
        if predictions.dim() > 1:
            predictions = predictions[0]  # Remove batch dimension
            
        # Get indices of positive predictions
        indices = torch.where(predictions > 0)[0].cpu().tolist()
        
        # Group by category
        result = {}
        for idx in indices:
            tag_name, category = self.dataset.get_tag_info(idx)
            
            if category not in result:
                result[category] = []
            
            if include_probabilities:
                prob = predictions[idx].item()
                result[category].append((tag_name, prob))
            else:
                result[category].append(tag_name)
        
        # Sort tags by probability within each category
        if include_probabilities:
            for category in result:
                result[category] = sorted(result[category], key=lambda x: x[1], reverse=True)
        
        return result

class FlashAttention(nn.Module):
    def __init__(self, dim, num_heads=8, dropout=0.1, batch_first=True):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.batch_first = batch_first
        self.head_dim = dim // num_heads
        assert self.head_dim * num_heads == dim, "dim must be divisible by num_heads"
        
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        
        for proj in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
            nn.init.xavier_uniform_(proj.weight, gain=0.1)
        
        self.scale = self.head_dim ** -0.5
        self.debug = False

    def _debug_print(self, name, tensor):
        """Debug helper"""
        if self.debug:
            print(f"\n{name}:")
            print(f"Shape: {tensor.shape}")
            print(f"Device: {tensor.device}")
            print(f"Dtype: {tensor.dtype}")
            if tensor.is_floating_point():
                with torch.no_grad():
                    print(f"Range: [{tensor.min().item():.3f}, {tensor.max().item():.3f}]")
                    print(f"Mean: {tensor.mean().item():.3f}")
                    print(f"Std: {tensor.std().item():.3f}")
    
    def _reshape_for_flash(self, x: torch.Tensor) -> torch.Tensor:
        """Reshape input tensor for flash attention format"""
        batch_size, seq_len, _ = x.size()
        x = x.view(batch_size, seq_len, self.num_heads, self.head_dim)
        x = x.transpose(1, 2)  # [B, H, S, D]
        return x.contiguous()

    def forward(self, query: torch.Tensor, key: Optional[torch.Tensor] = None, 
                value: Optional[torch.Tensor] = None, 
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Forward pass with flash attention"""
        if self.debug:
            print("\nFlashAttention Forward Pass")
        
        batch_size = query.size(0)
        
        # Use query as key/value if not provided
        key = query if key is None else key
        value = query if value is None else value
        
        # Project inputs
        q = self.q_proj(query)
        k = self.k_proj(key)
        v = self.v_proj(value)
        
        if self.debug:
            self._debug_print("Query before reshape", q)
        
        # Reshape for attention [B, H, S, D]
        q = self._reshape_for_flash(q)
        k = self._reshape_for_flash(k)
        v = self._reshape_for_flash(v)
        
        if self.debug:
            self._debug_print("Query after reshape", q)
        
        # Handle masking
        if mask is not None:
            # First convert mask to proper shape based on input dimensionality
            if mask.dim() == 2:  # [B, S]
                mask = mask.view(batch_size, 1, -1, 1)
            elif mask.dim() == 3:  # [B, S, S]
                mask = mask.view(batch_size, 1, mask.size(1), mask.size(2))
            elif mask.dim() == 5:  # [B, 1, S, S, S]
                mask = mask.squeeze(1).view(batch_size, 1, mask.size(2), mask.size(3))
                
            # Ensure mask is float16 if we're using float16
            mask = mask.to(q.dtype)
            
            if self.debug:
                self._debug_print("Prepared mask", mask)
                print(f"q shape: {q.shape}, mask shape: {mask.shape}")
            
            # Create attention mask that covers the full sequence length
            seq_len = q.size(2)
            if mask.size(-1) != seq_len:
                # Pad or trim mask to match sequence length
                new_mask = torch.zeros(batch_size, 1, seq_len, seq_len,
                                    device=mask.device, dtype=mask.dtype)
                min_len = min(seq_len, mask.size(-1))
                new_mask[..., :min_len, :min_len] = mask[..., :min_len, :min_len]
                mask = new_mask
            
            # Create key padding mask
            key_padding_mask = mask.squeeze(1).sum(-1) > 0
            key_padding_mask = key_padding_mask.view(batch_size, 1, -1, 1)
            
            # Apply the key padding mask
            k = k * key_padding_mask
            v = v * key_padding_mask
        
        if self.debug:
            self._debug_print("Query before attention", q)
            self._debug_print("Key before attention", k)
            self._debug_print("Value before attention", v)
        
        # Run flash attention
        dropout_p = self.dropout if self.training else 0.0
        output = flash_attn_func(
            q, k, v,
            dropout_p=dropout_p,
            softmax_scale=self.scale,
            causal=False
        )
        
        if self.debug:
            self._debug_print("Output after attention", output)
        
        # Reshape output [B, H, S, D] -> [B, S, H, D] -> [B, S, D]
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, -1, self.dim)
        
        # Final projection
        output = self.out_proj(output)
        
        if self.debug:
            self._debug_print("Final output", output)
        
        return output
    
class OptimizedTagEmbedding(nn.Module):
    def __init__(self, num_tags, embedding_dim, num_heads=8, dropout=0.1):
        super().__init__()
        # Single shared embedding for all tags
        self.embedding = nn.Embedding(num_tags, embedding_dim)
        self.attention = FlashAttention(embedding_dim, num_heads, dropout)
        self.norm1 = nn.LayerNorm(embedding_dim)
        self.norm2 = nn.LayerNorm(embedding_dim)
        
        # Single importance weighting for all tags
        self.tag_importance = nn.Parameter(torch.ones(num_tags) * 0.1)
        
        # Projection layers for unified tag context
        self.context_proj = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim * 2),
            nn.LayerNorm(embedding_dim * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim * 2, embedding_dim),
            nn.LayerNorm(embedding_dim)
        )
        
        self.importance_scale = nn.Parameter(torch.tensor(0.1))
        self.context_scale = nn.Parameter(torch.tensor(1.0))
        self.debug = False
        
    def _debug_print(self, name, tensor, extra_info=None):
        """Memory efficient debug printing with type handling"""
        if self.debug:
            print(f"\n{name}:")
            print(f"- Shape: {tensor.shape}")
            if isinstance(tensor, torch.Tensor):
                with torch.no_grad():
                    print(f"- Device: {tensor.device}")
                    print(f"- Dtype: {tensor.dtype}")
                    
                    # Convert to float32 for statistics if needed
                    if tensor.dtype not in [torch.float16, torch.float32, torch.float64]:
                        calc_tensor = tensor.float()
                    else:
                        calc_tensor = tensor
                        
                    try:
                        min_val = calc_tensor.min().item()
                        max_val = calc_tensor.max().item()
                        mean_val = calc_tensor.mean().item()
                        std_val = calc_tensor.std().item()
                        norm_val = torch.norm(calc_tensor).item()
                        
                        print(f"- Value range: [{min_val:.3f}, {max_val:.3f}]")
                        print(f"- Mean: {mean_val:.3f}")
                        print(f"- Std: {std_val:.3f}")
                        print(f"- L2 Norm: {norm_val:.3f}")
                        
                        if extra_info:
                            print(f"- Additional info: {extra_info}")
                    except Exception as e:
                        print(f"- Could not compute statistics: {str(e)}")
    
    def _debug_tensor(self, name, tensor):
        """Debug helper with dtype-specific analysis"""
        if self.debug and isinstance(tensor, torch.Tensor):
            print(f"\n{name}:")
            print(f"- Shape: {tensor.shape}")
            print(f"- Device: {tensor.device}")
            print(f"- Dtype: {tensor.dtype}")
            with torch.no_grad():
                has_nan = torch.isnan(tensor).any().item() if tensor.is_floating_point() else False
                has_inf = torch.isinf(tensor).any().item() if tensor.is_floating_point() else False
                print(f"- Contains NaN: {has_nan}")
                print(f"- Contains Inf: {has_inf}")
                
                # Different stats for different dtypes
                if tensor.is_floating_point():
                    print(f"- Range: [{tensor.min().item():.3f}, {tensor.max().item():.3f}]")
                    print(f"- Mean: {tensor.mean().item():.3f}")
                    print(f"- Std: {tensor.std().item():.3f}")
                else:
                    # For integer tensors
                    print(f"- Range: [{tensor.min().item()}, {tensor.max().item()}]")
                    print(f"- Unique values: {tensor.unique().numel()}")
    
    def _process_category(self, indices, masks):
        """Process a single category of tags"""
        # Get embeddings for this category
        embeddings = self.embedding(indices)
        
        if self.debug:
            self._debug_tensor("Category embeddings", embeddings)
        
        # Apply importance weights
        importance = torch.sigmoid(self.tag_importance) * self.importance_scale
        importance = torch.clamp(importance, min=0.01, max=10.0)
        importance_weights = importance[indices].unsqueeze(-1)
        
        # Apply and normalize
        embeddings = embeddings * importance_weights
        embeddings = self.norm1(embeddings)
        
        # Apply attention if we have more than one tag
        if embeddings.size(1) > 1:
            if masks is not None:
                attention_mask = torch.einsum('bi,bj->bij', masks, masks)
                attended = self.attention(embeddings, mask=attention_mask)
            else:
                attended = self.attention(embeddings)
            embeddings = self.norm2(attended)
        
        # Pool embeddings with masking
        if masks is not None:
            masked_embeddings = embeddings * masks.unsqueeze(-1)
            pooled = masked_embeddings.sum(dim=1) / masks.sum(dim=1, keepdim=True).clamp(min=1.0)
        else:
            pooled = embeddings.mean(dim=1)
            
        return pooled, embeddings
    
    def forward(self, tag_indices_dict, tag_masks_dict=None):
        """
        Process all tags in a unified embedding space
        Args:
            tag_indices_dict: dict of {category: tensor of indices}
            tag_masks_dict: dict of {category: tensor of masks}
        """
        if self.debug:
            print("\nOptimizedTagEmbedding Forward Pass")
        
        # Concatenate all indices and masks
        all_indices = []
        all_masks = []
        batch_size = None
        
        for category, indices in tag_indices_dict.items():
            if batch_size is None:
                batch_size = indices.size(0)
            all_indices.append(indices)
            if tag_masks_dict:
                all_masks.append(tag_masks_dict[category])
        
        # Stack along sequence dimension
        combined_indices = torch.cat(all_indices, dim=1)  # [B, total_seq_len]
        if tag_masks_dict:
            combined_masks = torch.cat(all_masks, dim=1)  # [B, total_seq_len]
        
        if self.debug:
            self._debug_tensor("Combined indices", combined_indices)
            if tag_masks_dict:
                self._debug_tensor("Combined masks", combined_masks)
        
        # Get embeddings for all tags using shared embedding
        embeddings = self.embedding(combined_indices)  # [B, total_seq_len, D]
        
        if self.debug:
            self._debug_tensor("Base embeddings", embeddings)
        
        # Apply unified importance weighting
        importance = torch.sigmoid(self.tag_importance) * self.importance_scale
        importance = torch.clamp(importance, min=0.01, max=10.0)
        importance_weights = importance[combined_indices].unsqueeze(-1)
        
        # Apply and normalize importance weights
        embeddings = embeddings * importance_weights
        embeddings = self.norm1(embeddings)
        
        if self.debug:
            self._debug_tensor("Weighted embeddings", embeddings)
        
        # Apply attention across all tags together
        if tag_masks_dict:
            attention_mask = torch.einsum('bi,bj->bij', combined_masks, combined_masks)
            attended = self.attention(embeddings, mask=attention_mask)
        else:
            attended = self.attention(embeddings)
        
        attended = self.norm2(attended)
        
        if self.debug:
            self._debug_tensor("Attended embeddings", attended)
        
        # Global pooling with masking
        if tag_masks_dict:
            masked_embeddings = attended * combined_masks.unsqueeze(-1)
            tag_context = masked_embeddings.sum(dim=1) / combined_masks.sum(dim=1, keepdim=True).clamp(min=1.0)
        else:
            tag_context = attended.mean(dim=1)
        
        # Project and scale context
        tag_context = self.context_proj(tag_context)
        context_scale = torch.clamp(self.context_scale, min=0.1, max=10.0)
        tag_context = tag_context * context_scale
        
        if self.debug:
            self._debug_tensor("Final tag context", tag_context)
        
        return tag_context, attended

class TagDataset:
    """Lightweight dataset wrapper for inference only"""
    def __init__(self, total_tags, idx_to_tag, tag_to_category):
        self.total_tags = total_tags
        self.idx_to_tag = idx_to_tag if isinstance(idx_to_tag, dict) else {int(k): v for k, v in idx_to_tag.items()}
        self.tag_to_category = tag_to_category
    
    def get_tag_info(self, idx):
        """Get tag name and category for a given index"""
        tag_name = self.idx_to_tag.get(idx, f"unknown-{idx}")
        category = self.tag_to_category.get(tag_name, "general")
        return tag_name, category

class ImageTagger(nn.Module):
    def __init__(self, total_tags, dataset, model_name='efficientnet_v2_l',
                 num_heads=16, dropout=0.1, pretrained=True,
                 tag_context_size=256):
        super().__init__()
        # Debug and stats flags
        self._flags = {
            'debug': False,
            'model_stats': False
        }
        
        # Core model config
        self.dataset = dataset
        self.tag_context_size = tag_context_size
        self.embedding_dim = 1280  # Fixed to EfficientNetV2-L output dimension
        
        # Initialize backbone
        if model_name == 'efficientnet_v2_l':
            weights = EfficientNet_V2_L_Weights.DEFAULT if pretrained else None
            self.backbone = efficientnet_v2_l(weights=weights)
            self.backbone.classifier = nn.Identity()
        
        # Spatial pooling only - no projection
        self.spatial_pool = nn.AdaptiveAvgPool2d((1, 1))
        
        # Initial tag prediction with bottleneck
        self.initial_classifier = nn.Sequential(
            nn.Linear(self.embedding_dim, self.embedding_dim * 2),
            nn.LayerNorm(self.embedding_dim * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(self.embedding_dim * 2, self.embedding_dim),
            nn.LayerNorm(self.embedding_dim),
            nn.GELU(),
            nn.Linear(self.embedding_dim, total_tags)
        )
        
        # Tag embeddings at full dimension
        self.tag_embedding = nn.Embedding(total_tags, self.embedding_dim)
        self.tag_attention = FlashAttention(self.embedding_dim, num_heads, dropout)
        self.tag_norm = nn.LayerNorm(self.embedding_dim)

        # Improved cross attention projection
        self.cross_proj = nn.Sequential(
            nn.Linear(self.embedding_dim, self.embedding_dim * 2),
            nn.LayerNorm(self.embedding_dim * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(self.embedding_dim * 2, self.embedding_dim)
        )
        
        # Cross attention at full dimension
        self.cross_attention = FlashAttention(self.embedding_dim, num_heads, dropout)
        self.cross_norm = nn.LayerNorm(self.embedding_dim)

        # Refined classifier with improved bottleneck
        self.refined_classifier = nn.Sequential(
            nn.Linear(self.embedding_dim * 2, self.embedding_dim * 2),  # Doubled input size for residual
            nn.LayerNorm(self.embedding_dim * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(self.embedding_dim * 2, self.embedding_dim),
            nn.LayerNorm(self.embedding_dim),
            nn.GELU(),
            nn.Linear(self.embedding_dim, total_tags)
        )
        
        # Temperature scaling
        self.temperature = nn.Parameter(torch.ones(1) * 1.5)
        
    def _get_selected_tags(self, logits):
        """Select top-K tags based on prediction confidence"""
        # Apply sigmoid to get probabilities
        probs = torch.sigmoid(logits)
        
        # Get top-K predictions for each image in batch
        batch_size = logits.size(0)
        topk_values, topk_indices = torch.topk(
            probs, k=self.tag_context_size, dim=1, largest=True, sorted=True
        )
        
        return topk_indices, topk_values
    
    @property
    def debug(self):
        return self._flags['debug']
    
    @debug.setter
    def debug(self, value):
        self._flags['debug'] = value
        
    @property
    def model_stats(self):
        return self._flags['model_stats']
    
    @model_stats.setter
    def model_stats(self, value):
        self._flags['model_stats'] = value
        
    def preprocess_image(self, image_path, image_size=512):
        """Process an image for inference using same preprocessing as training"""
        if not os.path.exists(image_path):
            raise ValueError(f"Image not found at path: {image_path}")
       
        # Initialize the same transform used during training
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
       
        try:
            with Image.open(image_path) as img:
                # Convert RGBA or Palette images to RGB
                if img.mode in ('RGBA', 'P'):
                    img = img.convert('RGB')
               
                # Get original dimensions
                width, height = img.size
                aspect_ratio = width / height
               
                # Calculate new dimensions to maintain aspect ratio
                if aspect_ratio > 1:
                    new_width = image_size
                    new_height = int(new_width / aspect_ratio)
                else:
                    new_height = image_size
                    new_width = int(new_height * aspect_ratio)
               
                # Resize with LANCZOS filter
                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
               
                # Create new image with padding
                new_image = Image.new('RGB', (image_size, image_size), (0, 0, 0))
                paste_x = (image_size - new_width) // 2
                paste_y = (image_size - new_height) // 2
                new_image.paste(img, (paste_x, paste_y))
               
                # Apply transforms (without normalization)
                img_tensor = transform(new_image)
                return img_tensor
        except Exception as e:
            raise Exception(f"Error processing {image_path}: {str(e)}")
    
    def forward(self, x):
        """Forward pass with simplified feature handling"""
        # Initialize tracking dicts
        model_stats = {} if self.model_stats else {}
        debug_tensors = {} if self.debug else None
        
        # 1. Image Feature Extraction
        features = self.backbone.features(x)
        features = self.spatial_pool(features).squeeze(-1).squeeze(-1)
        
        # 2. Initial Tag Predictions
        initial_logits = self.initial_classifier(features)
        initial_preds = torch.clamp(initial_logits / self.temperature, min=-15.0, max=15.0)
        
        # 3. Tag Selection & Embedding (simplified)
        pred_tag_indices, _ = self._get_selected_tags(initial_preds)
        tag_embeddings = self.tag_embedding(pred_tag_indices)
        
        # 4. Self-Attention on Tags
        attended_tags = self.tag_attention(tag_embeddings)
        attended_tags = self.tag_norm(attended_tags)
        
        # 5. Cross-Attention between Features and Tags
        features_proj = self.cross_proj(features)
        features_expanded = features_proj.unsqueeze(1).expand(-1, self.tag_context_size, -1)
        
        cross_attended = self.cross_attention(features_expanded, attended_tags)
        cross_attended = self.cross_norm(cross_attended)
        
        # 6. Feature Fusion with Residual Connection
        fused_features = cross_attended.mean(dim=1)  # Average across tag dimension
        # Concatenate original and attended features
        combined_features = torch.cat([features, fused_features], dim=-1)
        
        # 7. Refined Predictions
        refined_logits = self.refined_classifier(combined_features)
        refined_preds = torch.clamp(refined_logits / self.temperature, min=-15.0, max=15.0)
        
        # Return both prediction sets
        return initial_preds, refined_preds
    
    def predict(self, image_path, threshold=0.325, category_thresholds=None):
        """
        Run inference on an image with support for category-specific thresholds.
        """
        # Preprocess the image
        img_tensor = self.preprocess_image(image_path).unsqueeze(0)
        
        # Move to the same device as model and convert to half precision
        device = next(self.parameters()).device
        dtype = next(self.parameters()).dtype  # Match model's precision
        img_tensor = img_tensor.to(device, dtype=dtype)
        
        # Run inference
        with torch.no_grad():
            initial_preds, refined_preds = self.forward(img_tensor)
            
            # Apply sigmoid to get probabilities
            initial_probs = torch.sigmoid(initial_preds)
            refined_probs = torch.sigmoid(refined_preds)
            
            # Apply thresholds
            if category_thresholds:
                # Create binary prediction tensors
                refined_binary = torch.zeros_like(refined_probs)
                
                # Apply thresholds by category
                for category, cat_threshold in category_thresholds.items():
                    # Create a mask for tags in this category
                    category_mask = torch.zeros_like(refined_probs, dtype=torch.bool)
                    
                    # Find indices for this category
                    for tag_idx in range(refined_probs.size(-1)):
                        try:
                            _, tag_category = self.dataset.get_tag_info(tag_idx)
                            if tag_category == category:
                                category_mask[:, tag_idx] = True
                        except:
                            continue
                    
                    # Apply threshold only to tags in this category - ensure dtype consistency
                    cat_threshold_tensor = torch.tensor(cat_threshold, device=device, dtype=dtype)
                    refined_binary[category_mask] = (refined_probs[category_mask] >= cat_threshold_tensor).to(dtype)
                
                predictions = refined_binary
            else:
                # Use the same threshold for all tags
                threshold_tensor = torch.tensor(threshold, device=device, dtype=dtype)
                predictions = (refined_probs >= threshold_tensor).to(dtype)
            
            # Return both probabilities and thresholded predictions
            return {
                'initial_probabilities': initial_probs,
                'refined_probabilities': refined_probs,
                'predictions': predictions
            }
    
    def get_tags_from_predictions(self, predictions, include_probabilities=True):
        """
        Convert model predictions to human-readable tags grouped by category.
        """
        # Get non-zero predictions
        if predictions.dim() > 1:
            predictions = predictions[0]  # Remove batch dimension
            
        # Get indices of positive predictions
        indices = torch.where(predictions > 0)[0].cpu().tolist()
        
        # Group by category
        result = {}
        for idx in indices:
            tag_name, category = self.dataset.get_tag_info(idx)
            
            if category not in result:
                result[category] = []
            
            if include_probabilities:
                prob = predictions[idx].item()
                result[category].append((tag_name, prob))
            else:
                result[category].append(tag_name)
        
        # Sort tags by probability within each category
        if include_probabilities:
            for category in result:
                result[category] = sorted(result[category], key=lambda x: x[1], reverse=True)
        
        return result

def load_model(model_dir, device='cuda'):
    """Load model with better error handling and warnings"""
    print(f"Loading model from {model_dir}")
    
    try:
        # Load metadata
        metadata_path = os.path.join(model_dir, "metadata.json")
        if not os.path.exists(metadata_path):
            raise FileNotFoundError(f"Metadata file not found at {metadata_path}")
            
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
        
        # Load model info
        model_info_path = os.path.join(model_dir, "model_info_initial_only.json") 
        if os.path.exists(model_info_path):
            with open(model_info_path, 'r') as f:
                model_info = json.load(f)
        else:
            print("WARNING: Model info file not found, using default settings")
            model_info = {
                "tag_context_size": 256,
                "num_heads": 16,
                "precision": "float16"
            }
        
        # Create dataset wrapper
        dataset = TagDataset(
            total_tags=metadata['total_tags'],
            idx_to_tag=metadata['idx_to_tag'],
            tag_to_category=metadata['tag_to_category']
        )
        
        # Initialize model with exact settings from model_info
        model = ImageTagger(
            total_tags=metadata['total_tags'],
            dataset=dataset,
            num_heads=model_info.get('num_heads', 16),
            tag_context_size=model_info.get('tag_context_size', 256),
            pretrained=False
        )
        
        # Load weights
        state_dict_path = os.path.join(model_dir, "model.pt")
        if not os.path.exists(state_dict_path):
            raise FileNotFoundError(f"Model state dict not found at {state_dict_path}")
            
        state_dict = torch.load(state_dict_path, map_location=device)
        
        # First try strict loading
        try:
            model.load_state_dict(state_dict, strict=True)
            print("✓ Model state dict loaded with strict=True successfully")
        except Exception as e:
            print(f"! Strict loading failed: {str(e)}")
            print("Attempting non-strict loading...")
            
            # Try non-strict loading
            missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
            
            print(f"Non-strict loading completed with:")
            print(f"- {len(missing_keys)} missing keys")
            print(f"- {len(unexpected_keys)} unexpected keys")
            
            if len(missing_keys) > 0:
                print(f"Sample missing keys: {missing_keys[:5]}")
            if len(unexpected_keys) > 0:
                print(f"Sample unexpected keys: {unexpected_keys[:5]}")
        
        # Move model to device
        model = model.to(device)
        
        # Set to half precision if needed
        if model_info.get('precision') == 'float16':
            model = model.half()
            print("✓ Model converted to half precision")
        
        # Set to eval mode
        model.eval()
        print("✓ Model set to evaluation mode")
        
        # Verify parameter dtype
        param_dtype = next(model.parameters()).dtype
        print(f"✓ Model loaded with precision: {param_dtype}")
        
        return model, dataset
    
    except Exception as e:
        print(f"ERROR loading model: {str(e)}")
        import traceback
        traceback.print_exc()
        raise

# Example usage
if __name__ == "__main__":
    import sys
    
    # Get model directory from command line or use default
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "./exported_model"
    
    # Load model
    model, dataset, thresholds = load_model(model_dir)
    
    # Display info
    print(f"\nModel information:")
    print(f"  Total tags: {dataset.total_tags}")
    print(f"  Device: {next(model.parameters()).device}")
    print(f"  Precision: {next(model.parameters()).dtype}")
    
    # Test on an image if provided
    if len(sys.argv) > 2:
        image_path = sys.argv[2]
        print(f"\nRunning inference on {image_path}")
        
        # Use category thresholds if available
        if thresholds and 'categories' in thresholds:
            category_thresholds = {cat: opt['balanced']['threshold'] 
                              for cat, opt in thresholds['categories'].items()}
            results = model.predict(image_path, category_thresholds=category_thresholds)
        else:
            results = model.predict(image_path)
        
        # Get tags
        tags = model.get_tags_from_predictions(results['predictions'])
        
        # Print tags by category
        print("\nPredicted tags:")
        for category, category_tags in tags.items():
            print(f"\n{category.capitalize()}:")
            for tag, prob in category_tags:
                print(f"  {tag}: {prob:.3f}")