core/models.py

"""
Model management and optimization for BackgroundFX Pro.
Fixes MatAnyone quality issues and manages model loading.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Any, Optional, Tuple, List
from dataclasses import dataclass
import numpy as np
from pathlib import Path
import logging
import gc
from functools import lru_cache
import warnings

logger = logging.getLogger(__name__)


@dataclass
class ModelConfig:
    """Configuration for model management."""
    sam2_checkpoint: str = "checkpoints/sam2_hiera_large.pt"
    matanyone_checkpoint: str = "checkpoints/matanyone_v2.pth"
    device: str = "cuda"
    dtype: torch.dtype = torch.float16
    optimize_memory: bool = True
    use_amp: bool = True
    cache_size: int = 5
    enable_quality_fixes: bool = True
    matanyone_enhancement: bool = True
    use_tensorrt: bool = False
    batch_size: int = 1


class ModelCache:
    """Intelligent model caching system."""
    
    def __init__(self, max_size: int = 5):
        self.cache = {}
        self.max_size = max_size
        self.access_count = {}
        self.memory_usage = {}
        
    def add(self, key: str, model: Any, memory_size: float):
        """Add model to cache with memory tracking."""
        if len(self.cache) >= self.max_size:
            # Remove least recently used
            lru_key = min(self.access_count, key=self.access_count.get)
            self.remove(lru_key)
        
        self.cache[key] = model
        self.access_count[key] = 0
        self.memory_usage[key] = memory_size
        
    def get(self, key: str) -> Optional[Any]:
        """Get model from cache."""
        if key in self.cache:
            self.access_count[key] += 1
            return self.cache[key]
        return None
    
    def remove(self, key: str):
        """Remove model from cache and free memory."""
        if key in self.cache:
            model = self.cache[key]
            del self.cache[key]
            del self.access_count[key]
            del self.memory_usage[key]
            
            # Force cleanup
            del model
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    
    def clear(self):
        """Clear entire cache."""
        keys = list(self.cache.keys())
        for key in keys:
            self.remove(key)


class MatAnyoneModel(nn.Module):
    """Enhanced MatAnyone model with quality fixes."""
    
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.base_model = None
        self.quality_enhancer = QualityEnhancer() if config.enable_quality_fixes else None
        self.loaded = False
        
    def load(self):
        """Load MatAnyone model with optimizations."""
        if self.loaded:
            return
            
        try:
            # Load checkpoint
            checkpoint_path = Path(self.config.matanyone_checkpoint)
            if not checkpoint_path.exists():
                logger.warning(f"MatAnyone checkpoint not found at {checkpoint_path}")
                return
            
            # Load model weights
            state_dict = torch.load(
                checkpoint_path,
                map_location=self.config.device
            )
            
            # Initialize base model (placeholder - replace with actual MatAnyone architecture)
            self.base_model = self._build_matanyone_architecture()
            
            # Load weights with compatibility fixes
            self._load_weights_safe(state_dict)
            
            # Optimize model
            if self.config.optimize_memory:
                self._optimize_model()
            
            self.loaded = True
            logger.info("MatAnyone model loaded successfully")
            
        except Exception as e:
            logger.error(f"Failed to load MatAnyone model: {e}")
            self.loaded = False
    
    def _build_matanyone_architecture(self) -> nn.Module:
        """Build MatAnyone architecture."""
        # This is a placeholder - replace with actual MatAnyone architecture
        class MatAnyoneBase(nn.Module):
            def __init__(self):
                super().__init__()
                self.encoder = nn.Sequential(
                    nn.Conv2d(4, 64, 3, padding=1),
                    nn.ReLU(),
                    nn.Conv2d(64, 128, 3, stride=2, padding=1),
                    nn.ReLU(),
                    nn.Conv2d(128, 256, 3, stride=2, padding=1),
                    nn.ReLU(),
                )
                self.decoder = nn.Sequential(
                    nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),
                    nn.ReLU(),
                    nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),
                    nn.ReLU(),
                    nn.Conv2d(64, 4, 3, padding=1),
                    nn.Sigmoid()
                )
                
            def forward(self, x):
                features = self.encoder(x)
                output = self.decoder(features)
                return output
        
        return MatAnyoneBase().to(self.config.device)
    
    def _load_weights_safe(self, state_dict: Dict):
        """Safely load weights with compatibility handling."""
        model_dict = self.base_model.state_dict()
        
        # Filter compatible weights
        compatible_dict = {}
        for k, v in state_dict.items():
            # Remove module prefix if present
            if k.startswith('module.'):
                k = k[7:]
            
            if k in model_dict and model_dict[k].shape == v.shape:
                compatible_dict[k] = v
            else:
                logger.warning(f"Skipping incompatible weight: {k}")
        
        # Load compatible weights
        model_dict.update(compatible_dict)
        self.base_model.load_state_dict(model_dict, strict=False)
        
        logger.info(f"Loaded {len(compatible_dict)}/{len(state_dict)} weights")
    
    def _optimize_model(self):
        """Optimize model for inference."""
        if not self.base_model:
            return
            
        self.base_model.eval()
        
        # Convert to half precision if using GPU
        if self.config.dtype == torch.float16 and self.config.device != "cpu":
            self.base_model = self.base_model.half()
        
        # Disable gradient computation
        for param in self.base_model.parameters():
            param.requires_grad = False
        
        # TensorRT optimization (if available)
        if self.config.use_tensorrt:
            try:
                self._optimize_with_tensorrt()
            except Exception as e:
                logger.warning(f"TensorRT optimization failed: {e}")
    
    def forward(self, image: torch.Tensor, mask: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Enhanced forward pass with quality fixes."""
        if not self.loaded:
            self.load()
        
        if not self.base_model:
            return {'alpha': mask, 'foreground': image}
        
        # Prepare input
        x = torch.cat([image, mask.unsqueeze(1)], dim=1)
        
        # Fix input quality issues
        if self.config.matanyone_enhancement:
            x = self._preprocess_input(x)
        
        # Forward pass with mixed precision
        with torch.cuda.amp.autocast(enabled=self.config.use_amp):
            output = self.base_model(x)
        
        # Parse output
        alpha = output[:, 3:4, :, :]
        foreground = output[:, :3, :, :]
        
        # Apply quality enhancement
        if self.quality_enhancer:
            alpha = self.quality_enhancer.enhance_alpha(alpha, mask)
            foreground = self.quality_enhancer.enhance_foreground(foreground, image)
        
        # Post-process to fix common MatAnyone issues
        alpha = self._fix_matanyone_artifacts(alpha, mask)
        
        return {
            'alpha': alpha,
            'foreground': foreground,
            'confidence': self._compute_confidence(alpha, mask)
        }
    
    def _preprocess_input(self, x: torch.Tensor) -> torch.Tensor:
        """Preprocess input to improve MatAnyone quality."""
        # Denoise input
        if x.shape[2] > 64:  # Only for reasonable resolutions
            x = self._bilateral_filter_torch(x)
        
        # Normalize properly
        x = torch.clamp(x, 0, 1)
        
        # Enhance edges in mask channel
        mask_channel = x[:, 3:4, :, :]
        mask_enhanced = self._enhance_mask_edges(mask_channel)
        x = torch.cat([x[:, :3, :, :], mask_enhanced], dim=1)
        
        return x
    
    def _fix_matanyone_artifacts(self, alpha: torch.Tensor,
                                 original_mask: torch.Tensor) -> torch.Tensor:
        """Fix common MatAnyone artifacts."""
        # Fix edge bleeding
        alpha = self._fix_edge_bleeding(alpha, original_mask)
        
        # Fix transparency issues
        alpha = self._fix_transparency_issues(alpha)
        
        # Ensure consistency with original mask
        alpha = self._ensure_mask_consistency(alpha, original_mask)
        
        return alpha
    
    def _fix_edge_bleeding(self, alpha: torch.Tensor,
                          original_mask: torch.Tensor) -> torch.Tensor:
        """Fix edge bleeding artifacts."""
        # Detect edges
        edges = self._detect_edges_torch(original_mask)
        
        # Create edge mask
        edge_mask = F.max_pool2d(edges, kernel_size=5, stride=1, padding=2)
        
        # Refine alpha near edges
        alpha_refined = alpha.clone()
        edge_region = edge_mask > 0.1
        
        # Apply guided filter near edges
        if edge_region.any():
            alpha_refined[edge_region] = (
                0.7 * alpha[edge_region] + 
                0.3 * original_mask.unsqueeze(1).expand_as(alpha)[edge_region]
            )
        
        return alpha_refined
    
    def _fix_transparency_issues(self, alpha: torch.Tensor) -> torch.Tensor:
        """Fix transparency artifacts."""
        # Identify problematic transparency values
        mid_range = (alpha > 0.2) & (alpha < 0.8)
        
        # Push mid-range values toward 0 or 1
        alpha_fixed = alpha.clone()
        alpha_fixed[mid_range] = torch.where(
            alpha[mid_range] > 0.5,
            torch.clamp(alpha[mid_range] * 1.2, max=1.0),
            torch.clamp(alpha[mid_range] * 0.8, min=0.0)
        )
        
        # Smooth transitions
        alpha_fixed = F.gaussian_blur(alpha_fixed, kernel_size=(3, 3))
        
        return alpha_fixed
    
    def _ensure_mask_consistency(self, alpha: torch.Tensor,
                                original_mask: torch.Tensor) -> torch.Tensor:
        """Ensure consistency with original mask."""
        # Expand mask dimensions if needed
        if original_mask.dim() == 2:
            original_mask = original_mask.unsqueeze(0).unsqueeze(0)
        elif original_mask.dim() == 3:
            original_mask = original_mask.unsqueeze(1)
        
        # Where original mask is 0, alpha should also be 0
        alpha = torch.where(original_mask < 0.1, torch.zeros_like(alpha), alpha)
        
        # Where original mask is 1, alpha should be close to 1
        alpha = torch.where(original_mask > 0.9, torch.ones_like(alpha) * 0.95, alpha)
        
        return alpha
    
    def _compute_confidence(self, alpha: torch.Tensor,
                          original_mask: torch.Tensor) -> torch.Tensor:
        """Compute confidence score for the output."""
        # Expand dimensions if needed
        if original_mask.dim() < alpha.dim():
            original_mask = original_mask.unsqueeze(1).expand_as(alpha)
        
        # Compute similarity
        diff = torch.abs(alpha - original_mask)
        confidence = 1.0 - torch.mean(diff, dim=(1, 2, 3))
        
        return confidence
    
    def _bilateral_filter_torch(self, x: torch.Tensor) -> torch.Tensor:
        """Apply bilateral filter in PyTorch."""
        # Simple approximation using Gaussian blur
        # For true bilateral filtering, would need custom CUDA kernel
        return F.gaussian_blur(x, kernel_size=(5, 5))
    
    def _enhance_mask_edges(self, mask: torch.Tensor) -> torch.Tensor:
        """Enhance edges in mask channel."""
        # Detect edges
        edges = self._detect_edges_torch(mask)
        
        # Enhance mask with edges
        enhanced = mask + 0.3 * edges
        enhanced = torch.clamp(enhanced, 0, 1)
        
        return enhanced
    
    def _detect_edges_torch(self, x: torch.Tensor) -> torch.Tensor:
        """Detect edges using Sobel filters."""
        # Sobel kernels
        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], 
                               dtype=x.dtype, device=x.device).view(1, 1, 3, 3)
        sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], 
                               dtype=x.dtype, device=x.device).view(1, 1, 3, 3)
        
        # Apply Sobel filters
        edges_x = F.conv2d(x, sobel_x, padding=1)
        edges_y = F.conv2d(x, sobel_y, padding=1)
        
        # Compute edge magnitude
        edges = torch.sqrt(edges_x ** 2 + edges_y ** 2)
        
        return edges


class SAM2Model:
    """SAM2 model wrapper with optimizations."""
    
    def __init__(self, config: ModelConfig):
        self.config = config
        self.model = None
        self.predictor = None
        self.loaded = False
        
    def load(self):
        """Load SAM2 model."""
        if self.loaded:
            return
            
        try:
            # Import SAM2 (assuming it's installed)
            from sam2.build_sam import build_sam2
            from sam2.sam2_image_predictor import SAM2ImagePredictor
            
            # Build model
            self.model = build_sam2(
                config_file="sam2_hiera_l.yaml",
                ckpt_path=self.config.sam2_checkpoint,
                device=self.config.device
            )
            
            # Create predictor
            self.predictor = SAM2ImagePredictor(self.model)
            
            self.loaded = True
            logger.info("SAM2 model loaded successfully")
            
        except Exception as e:
            logger.error(f"Failed to load SAM2 model: {e}")
            self.loaded = False
    
    def predict(self, image: np.ndarray, prompts: Optional[Dict] = None) -> np.ndarray:
        """Generate segmentation mask."""
        if not self.loaded:
            self.load()
        
        if not self.predictor:
            return np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
        
        # Set image
        self.predictor.set_image(image)
        
        # Use prompts if provided, otherwise use automatic segmentation
        if prompts:
            masks, scores, _ = self.predictor.predict(
                point_coords=prompts.get('points'),
                point_labels=prompts.get('labels'),
                box=prompts.get('box'),
                multimask_output=True
            )
            # Select best mask
            mask = masks[np.argmax(scores)]
        else:
            # Automatic segmentation
            masks = self.predictor.generate_auto_masks(image)
            mask = masks[0] if len(masks) > 0 else np.zeros_like(image[:, :, 0])
        
        return mask


class QualityEnhancer(nn.Module):
    """Neural quality enhancement module."""
    
    def __init__(self):
        super().__init__()
        self.alpha_refiner = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 16, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 1, 3, padding=1),
            nn.Sigmoid()
        )
        
        self.foreground_enhancer = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 3, 3, padding=1),
            nn.Tanh()
        )
    
    def enhance_alpha(self, alpha: torch.Tensor,
                     original_mask: torch.Tensor) -> torch.Tensor:
        """Enhance alpha channel quality."""
        # Refine with neural network
        refined = self.alpha_refiner(alpha)
        
        # Blend with original for stability
        enhanced = 0.7 * refined + 0.3 * alpha
        
        return torch.clamp(enhanced, 0, 1)
    
    def enhance_foreground(self, foreground: torch.Tensor,
                          original_image: torch.Tensor) -> torch.Tensor:
        """Enhance foreground quality."""
        # Compute residual
        residual = self.foreground_enhancer(foreground)
        
        # Add residual
        enhanced = foreground + 0.1 * residual
        
        return torch.clamp(enhanced, 0, 1)


class ModelManager:
    """Central model management system."""
    
    def __init__(self, config: Optional[ModelConfig] = None):
        self.config = config or ModelConfig()
        self.cache = ModelCache(max_size=self.config.cache_size)
        self.models = {}
        
        # Initialize models
        self.sam2 = SAM2Model(self.config)
        self.matanyone = MatAnyoneModel(self.config)
        
    def load_all(self):
        """Load all models."""
        logger.info("Loading all models...")
        self.sam2.load()
        self.matanyone.load()
        logger.info("All models loaded")
    
    def get_sam2(self) -> SAM2Model:
        """Get SAM2 model."""
        if not self.sam2.loaded:
            self.sam2.load()
        return self.sam2
    
    def get_matanyone(self) -> MatAnyoneModel:
        """Get MatAnyone model."""
        if not self.matanyone.loaded:
            self.matanyone.load()
        return self.matanyone
    
    def process_frame(self, image: np.ndarray,
                     mask: Optional[np.ndarray] = None) -> Dict[str, Any]:
        """Process single frame through pipeline."""
        # Convert to tensor
        image_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float() / 255.0
        image_tensor = image_tensor.to(self.config.device)
        
        # Get or generate mask
        if mask is None:
            mask = self.sam2.predict(image)
        
        mask_tensor = torch.from_numpy(mask).float().to(self.config.device)
        
        # Process with MatAnyone
        result = self.matanyone(image_tensor, mask_tensor)
        
        # Convert back to numpy
        output = {
            'alpha': result['alpha'].squeeze().cpu().numpy(),
            'foreground': result['foreground'].squeeze().permute(1, 2, 0).cpu().numpy() * 255,
            'confidence': result['confidence'].cpu().numpy()
        }
        
        return output
    
    def cleanup(self):
        """Cleanup models and free memory."""
        self.cache.clear()
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


# Export classes
__all__ = [
    'ModelManager',
    'SAM2Model', 
    'MatAnyoneModel',
    'ModelConfig',
    'ModelCache',
    'QualityEnhancer'
]