# blip3o_fast.py
# Training: Qwen3 + Grounding DINO + SAM-2 for mask supervision
# Inference: Lightweight - no external components needed

from typing import List, Optional, Tuple, Union, Dict, Any
import re
import json

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    Qwen2Config,
    Qwen2Model,
    Qwen2ForCausalLM
)
from transformers.modeling_outputs import CausalLMOutputWithPast
from diffusers.training_utils import (
    compute_density_for_timestep_sampling,
    compute_loss_weighting_for_sd3
)
from diffusers.utils.torch_utils import randn_tensor

from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM


# ============================================================
# TRAINING ONLY: Qwen3 Client for Instruction Parsing
# ============================================================

class Qwen3InstructionParser:
    """Parses edit instructions using Qwen3 LLM. Used only during training."""

    def __init__(
            self,
            model_name: str = "Qwen/Qwen3-1.7B",
            device: str = "cuda",
            torch_dtype: torch.dtype = torch.float16
    ):
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
            device_map=device
        )
        self.model.eval()
        self._cache: Dict[str, Dict] = {}

    @torch.no_grad()
    def parse(self, instruction: str) -> Dict[str, Any]:
        if instruction in self._cache:
            return self._cache[instruction]

        prompt = self._build_prompt(instruction)
        messages = [{"role": "user", "content": prompt}]
        text = self.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
        )
        inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
        outputs = self.model.generate(
            **inputs, max_new_tokens=256, temperature=0.1,
            do_sample=False, pad_token_id=self.tokenizer.eos_token_id
        )
        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        parsed = self._parse_response(response)
        self._cache[instruction] = parsed
        return parsed

    def _build_prompt(self, instruction: str) -> str:
        return f"""You are an image editing instruction parser. Extract structured information.

Respond ONLY with valid JSON:
{{"operation": "<type>", "source_object": "<object or null>", "target_object": "<object or null>", "location": "<location or null>", "attributes": "<attributes or null>"}}

Operation types: remove, replace, add, extract, style, adjust, compose, action, other

Examples:
"Remove the red car" -> {{"operation": "remove", "source_object": "red car", "target_object": null, "location": null, "attributes": null}}
"Replace the dog with a cat" -> {{"operation": "replace", "source_object": "dog", "target_object": "cat", "location": null, "attributes": null}}
"Make the dress blue" -> {{"operation": "adjust", "source_object": "dress", "target_object": null, "location": null, "attributes": "blue"}}

Input: "{instruction}"
Output:"""

    def _parse_response(self, response: str) -> Dict[str, Any]:
        default = {"operation": "other", "source_object": None, "target_object": None, "location": None,
                   "attributes": None}
        try:
            parsed = json.loads(response.strip())
        except json.JSONDecodeError:
            match = re.search(r'\{[^{}]*\}', response, re.DOTALL)
            if match:
                try:
                    parsed = json.loads(match.group())
                except:
                    return default
            else:
                return default
        for key in default:
            if key not in parsed:
                parsed[key] = default[key]
        valid_ops = ["remove", "replace", "add", "extract", "style", "adjust", "compose", "action", "other"]
        if parsed["operation"] not in valid_ops:
            parsed["operation"] = "other"
        return parsed


class SAM3MaskGenerator:
    """
    Generates segmentation masks using SAM3.
    SAM3 natively supports text prompts - no Grounding DINO needed!
    """

    def __init__(self, device: str = "cuda"):
        self.device = device
        self._model = None
        self._processor = None

    def _load_model(self):
        """Lazy load SAM3 model."""
        if self._model is None:
            from sam3.model_builder import build_sam3_image_model
            from sam3.model.sam3_image_processor import Sam3Processor

            print("Loading SAM3...")
            self._model = build_sam3_image_model()
            self._processor = Sam3Processor(self._model)
            print("SAM3 loaded!")

    def _prepare_image(self, image):
        """Convert various image formats to PIL Image."""
        from PIL import Image as PILImage

        if isinstance(image, PILImage.Image):
            return image.convert("RGB")
        elif isinstance(image, torch.Tensor):
            if image.dim() == 4:
                image = image[0]

            if image.dtype in (torch.bfloat16, torch.float16):
                image = image.float()
            if image.shape[0] in [1, 3]:
                image_np = image.permute(1, 2, 0).cpu().numpy()
            else:
                image_np = image.cpu().numpy()
            if image_np.max() <= 1.0:
                image_np = (image_np * 255).astype(np.uint8)
            else:
                image_np = image_np.astype(np.uint8)
            return PILImage.fromarray(image_np).convert("RGB")
        elif isinstance(image, np.ndarray):
            if image.max() <= 1.0:
                image = (image * 255).astype(np.uint8)
            return PILImage.fromarray(image).convert("RGB")
        else:
            return PILImage.fromarray(np.array(image)).convert("RGB")

    @torch.no_grad()
    def generate_mask(
            self,
            image,
            parsed: Dict,
            detect_all: bool = False
    ) -> torch.Tensor:
        """
        Generate segmentation mask using SAM3 with text prompt.

        Args:
            image: Input image (PIL, tensor, or numpy)
            parsed: Parsed instruction dict with 'source_object', 'operation', etc.
            detect_all: Whether to return all instances

        Returns:
            mask: [1, H, W] binary mask tensor
        """
        self._load_model()
        # Convert image to PIL
        image_pil = self._prepare_image(image)
        W, H = image_pil.size

        # Build text prompt from parsed instruction
        text_prompt = self._build_text_prompt(parsed)

        if not text_prompt:
            if parsed.get("operation") == "style":
                return torch.ones(1, H, W)
            return torch.zeros(1, H, W)

        # Set image in SAM3
        inference_state = self._processor.set_image(image_pil)

        # Get segmentation with text prompt
        output = self._processor.set_text_prompt(
            state=inference_state,
            prompt=text_prompt
        )

        masks = output["masks"]  # List of masks
        scores = output["scores"]  # Confidence scores

        if masks is None or len(masks) == 0:
            return torch.zeros(1, H, W)

        # Convert masks to tensor
        if isinstance(masks, np.ndarray):
            masks = torch.from_numpy(masks)
        elif isinstance(masks, list):
            masks = torch.stack([torch.from_numpy(m) if isinstance(m, np.ndarray) else m for m in masks])

        if detect_all:
            # Combine all masks
            combined_mask = masks.float().max(dim=0)[0]
            return combined_mask.unsqueeze(0)
        else:
            # Return highest scoring mask
            if isinstance(scores, (list, np.ndarray)):
                scores = torch.tensor(scores)
            best_idx = scores.argmax()
            return masks[best_idx].unsqueeze(0).float()

    def _build_text_prompt(self, parsed: Dict) -> str:
        """Build SAM3 text prompt from parsed instruction."""
        operation = parsed.get("operation", "other")
        source = parsed.get("source_object")
        target = parsed.get("target_object")
        location = parsed.get("location")
        attributes = parsed.get("attributes")

        if operation in ["remove", "replace", "extract", "adjust", "action"]:
            # Need to find the source object
            if source:
                # Add attributes if available
                if attributes and operation == "adjust":
                    return source  # e.g., "dress" for "make the dress blue"
                return source
        elif operation == "add":
            # For add, find where to add (the context object)
            if source:
                return source  # e.g., "woman" for "put sunglasses on the woman"
            elif location:
                return location
        elif operation == "compose":
            if source:
                return source
        elif operation == "style":
            # Style affects whole image, return empty
            return ""

        return source or ""


class EditMaskGenerator:
    """
    Complete mask generation pipeline using Qwen3 + SAM3.

    Simplified from: Qwen3 → Grounding DINO → SAM-2
    To:              Qwen3 → SAM3 (native text support)
    """

    def __init__(
            self,
            qwen_model: str = "Qwen/Qwen3-1.7B",
            device: str = "cuda",
            enabled: bool = True
    ):
        self.device = device
        self.enabled = enabled

        if enabled:
            print("Initializing EditMaskGenerator with SAM3...")
            self.parser = Qwen3InstructionParser(model_name=qwen_model, device=device)
            self.segmenter = SAM3MaskGenerator(device=device)
            print("EditMaskGenerator ready!")
        else:
            self.parser = None
            self.segmenter = None

    @torch.no_grad()
    def generate(
            self,
            image,
            instruction: str,
            return_parsed: bool = False
    ):
        """Generate edit mask from image and instruction."""
        if not self.enabled:
            if isinstance(image, torch.Tensor):
                H, W = image.shape[-2:]
            else:
                H, W = np.array(image).shape[:2]
            mask = torch.zeros(1, H, W)
            return (mask, {"operation": "other"}) if return_parsed else mask

        # Step 1: Parse instruction with Qwen3
        parsed = self.parser.parse(instruction)

        # Step 2: Generate mask with SAM3 (native text prompt!)
        detect_all = "all" in instruction.lower()
        mask = self.segmenter.generate_mask(image, parsed, detect_all=detect_all)

        return (mask, parsed) if return_parsed else mask


# ============================================================
# Model Configuration
# ============================================================
class blip3oFastConfig(Qwen2Config):
    model_type = "llava_qwen2"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.latent_channels = kwargs.get("latent_channels", 32)

        # Conditioning
        self.use_spatial_conditioning = kwargs.get("use_spatial_conditioning", False)
        self.use_mask_conditioning = kwargs.get("use_mask_conditioning", True)
        self.use_operation_embedding = kwargs.get("use_operation_embedding", True)
        self.use_mask_predictor = kwargs.get("use_mask_predictor", True)
        self.mask_predictor_loss_weight = kwargs.get("mask_predictor_loss_weight", 0.5)

        # Dropout
        self.spatial_drop_prob = kwargs.get("spatial_drop_prob", 0.1)
        self.mask_drop_prob = kwargs.get("mask_drop_prob", 0.1)

        # Mask generator config (SIMPLIFIED - no Grounding DINO!)
        self.mask_generator_enabled = kwargs.get("mask_generator_enabled", True)
        self.qwen_model = kwargs.get("qwen_model", "Qwen/Qwen3-1.7B")


# ============================================================
# Mask Predictor: Learns to predict edit regions from LLM hidden states
# ============================================================

class BF16SafeLayerNorm(nn.Module):
    """LayerNorm that works correctly with BF16 and PEFT."""
    
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        # Explicitly initialize with proper values
        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.float32))
        self.bias = nn.Parameter(torch.zeros(hidden_size, dtype=torch.float32))
        self.eps = eps
        self.hidden_size = hidden_size
        
        # Force initialization
        self.reset_parameters()
    
    def reset_parameters(self):
        """Ensure weights are properly initialized."""
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Always compute normalization in float32 for stability
        input_dtype = x.dtype
        x_f32 = x.float()
        
        # Manual layer norm computation
        mean = x_f32.mean(dim=-1, keepdim=True)
        var = x_f32.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x_f32 - mean) / torch.sqrt(var + self.eps)
        
        # Apply weight and bias in float32
        output = x_norm * self.weight.float() + self.bias.float()
        
        # Convert back to original dtype
        return output.to(input_dtype)


class MaskPredictor(nn.Module):
    """
    Predicts edit mask from LLM hidden states.
    This is the KEY component that enables mask-free inference.

    The mask predictor learns to identify WHICH object needs to be edited
    based on the instruction (e.g., "remove the white dog") and the image
    understanding encoded in the LLM hidden states.

    Architecture:
    1. Extract instruction-relevant features using attention pooling
    2. Project to spatial features
    3. Decode to mask
    """

    def __init__(self, hidden_size: int, latent_channels: int, latent_size: int = 32):
        super().__init__()

        self.latent_size = latent_size
        self.hidden_size = hidden_size

        # Attention pooling to focus on instruction-relevant tokens
        # Instead of simple mean pooling, learn which tokens are important
        self.attention_pool = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 4),
            nn.Tanh(),
            nn.Linear(hidden_size // 4, 1),
        )

        # Layer norm for stability
        self.input_norm = BF16SafeLayerNorm(hidden_size)

        # Project pooled features to spatial representation
        intermediate_size = hidden_size // 2
        spatial_dim = latent_size * latent_size * 64
        
        self.hidden_proj = nn.Sequential(
            nn.Linear(hidden_size, intermediate_size),
            nn.LayerNorm(intermediate_size),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(intermediate_size, intermediate_size),
            nn.LayerNorm(intermediate_size),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(intermediate_size, spatial_dim),
        )

        # Upsample to mask with more capacity
        self.mask_decoder = nn.Sequential(
            nn.Conv2d(64, 256, 3, padding=1),
            nn.GroupNorm(32, 256),
            nn.GELU(),
            nn.Conv2d(256, 128, 3, padding=1),
            nn.GroupNorm(16, 128),
            nn.GELU(),
            nn.Conv2d(128, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.GELU(),
            nn.Conv2d(64, 1, 1),
        )

        self._init_weights()

    def _init_weights(self):
        """Initialize weights for stable training."""
        # Initialize attention pooling
        for module in self.attention_pool:
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight, gain=0.1)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
        
        # Initialize LayerNorm
        if hasattr(self, 'input_norm'):
            self.input_norm.reset_parameters()
        
        # Initialize projection layers
        for module in self.hidden_proj:
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight, gain=0.1)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)
        
        # Initialize conv layers
        for module in self.mask_decoder:
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.GroupNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)
        
        # Initialize final layer with small weights for stable start
        for module in reversed(list(self.mask_decoder)):
            if isinstance(module, nn.Conv2d):
                nn.init.normal_(module.weight, mean=0.0, std=0.01)
                nn.init.zeros_(module.bias)
                break

    def forward(self, hidden_states: torch.Tensor, return_logits: bool = False) -> torch.Tensor:
        """
        Args:
            hidden_states: [B, seq_len, hidden_size] from LLM
            return_logits: If True, return logits instead of probabilities

        Returns:
            mask: [B, 1, H, W] predicted edit mask
        """
        batch_size = hidden_states.shape[0]
        device = hidden_states.device
        
        # Check for NaN/Inf in input
        if torch.isnan(hidden_states).any() or torch.isinf(hidden_states).any():
            print("WARNING: NaN/Inf in hidden_states input to MaskPredictor")
            if return_logits:
                return torch.zeros(batch_size, 1, self.latent_size, self.latent_size, 
                                   device=device, dtype=torch.float32, requires_grad=True)
            return torch.full((batch_size, 1, self.latent_size, self.latent_size), 0.5,
                              device=device, dtype=torch.float32, requires_grad=True)
        
        # Normalize hidden states
        hidden_states = self.input_norm(hidden_states)
        
        if torch.isnan(hidden_states).any():
            print("WARNING: NaN after input_norm in MaskPredictor")
            if return_logits:
                return torch.zeros(batch_size, 1, self.latent_size, self.latent_size, 
                                   device=device, dtype=torch.float32, requires_grad=True)
            return torch.full((batch_size, 1, self.latent_size, self.latent_size), 0.5,
                              device=device, dtype=torch.float32, requires_grad=True)

        # Get dtype from first layer
        target_dtype = self.attention_pool[0].weight.dtype
        hidden_states = hidden_states.to(target_dtype)
        
        # Attention pooling: learn which tokens are important for mask prediction
        # [B, seq_len, hidden_size] -> [B, seq_len, 1]
        attn_weights = self.attention_pool(hidden_states)
        attn_weights = F.softmax(attn_weights, dim=1)  # [B, seq_len, 1]
        
        # Weighted sum of hidden states
        # [B, seq_len, hidden_size] * [B, seq_len, 1] -> [B, hidden_size]
        pooled = (hidden_states * attn_weights).sum(dim=1)

        # Project to spatial features
        spatial = self.hidden_proj(pooled)  # [B, spatial_dim]
        spatial = spatial.view(-1, 64, self.latent_size, self.latent_size)  # [B, 64, H, W]

        # Decode to mask logits
        mask_logits = self.mask_decoder(spatial)  # [B, 1, H, W]

        if return_logits:
            return mask_logits.float()

        # Apply sigmoid to get probabilities
        mask = torch.sigmoid(mask_logits.float())
        return mask


# ============================================================
# Main Model
# ============================================================

class blip3oFastModel(LlavaMetaModel, Qwen2Model):
    config_class = blip3oFastConfig

    def __init__(self, config: Qwen2Config):
        super(blip3oFastModel, self).__init__(config)


class blip3oFastForCausalLM(Qwen2ForCausalLM, LlavaMetaForCausalLM):
    config_class = blip3oFastConfig

    def __init__(self, config):
        super(Qwen2ForCausalLM, self).__init__(config)

        self.model = blip3oFastModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        latent_channels = getattr(config, 'latent_channels', 32)

        # ============================================================
        # Spatial Reference Encoder
        # ============================================================
        if getattr(config, 'use_spatial_conditioning', True):
            self.spatial_ref_encoder = nn.Sequential(
                nn.Conv2d(latent_channels, 320, 3, padding=1),
                nn.GroupNorm(32, 320),
                nn.SiLU(),
                nn.Conv2d(320, 320, 3, padding=1),
                nn.GroupNorm(32, 320),
                nn.SiLU(),
                nn.Conv2d(320, latent_channels, 3, padding=1),
            )
            self.spatial_weight = nn.Parameter(torch.tensor(0.0))
        else:
            self.spatial_ref_encoder = None
            self.spatial_weight = None

        # ============================================================
        # Mask Encoder (encodes mask into conditioning)
        # ============================================================
        if getattr(config, 'use_mask_conditioning', True):
            self.mask_encoder = nn.Sequential(
                nn.Conv2d(1, 64, 3, padding=1),
                nn.GroupNorm(8, 64),
                nn.SiLU(),
                nn.Conv2d(64, 128, 3, padding=1),
                nn.GroupNorm(16, 128),
                nn.SiLU(),
                nn.Conv2d(128, latent_channels, 3, padding=1),
            )
            self.mask_weight = nn.Parameter(torch.tensor(0.0))
        else:
            self.mask_encoder = None
            self.mask_weight = None

        # ============================================================
        # Mask Predictor (CRITICAL: enables mask-free inference)
        # ============================================================
        if getattr(config, 'use_mask_predictor', True):
            self.mask_predictor = MaskPredictor(
                hidden_size=config.hidden_size,
                latent_channels=latent_channels,
                latent_size=32  # Adjust based on your latent resolution
            )
        else:
            self.mask_predictor = None

        # ============================================================
        # Operation Embedding
        # ============================================================
        if getattr(config, 'use_operation_embedding', True):
            self.operation_types = ["remove", "replace", "add", "extract", "style", "adjust", "compose", "action",
                                    "other"]
            self.operation_embedding = nn.Embedding(len(self.operation_types), latent_channels)
        else:
            self.operation_types = None
            self.operation_embedding = None

        # Mask generator (training only, lazy init)
        self._mask_generator = None
        self._mask_generator_initialized = False

        self._init_conditioning_layers()
        self.post_init()

    def _init_conditioning_layers(self):
        """Initialize conditioning layers. Called during __init__ and can be called after loading."""
        if self.spatial_ref_encoder is not None:
            for module in self.spatial_ref_encoder:
                if isinstance(module, nn.Conv2d):
                    nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                    if module.bias is not None:
                        nn.init.zeros_(module.bias)
                elif isinstance(module, nn.GroupNorm):
                    nn.init.ones_(module.weight)
                    nn.init.zeros_(module.bias)
            # Zero-init the last layer
            nn.init.zeros_(self.spatial_ref_encoder[-1].weight)
            nn.init.zeros_(self.spatial_ref_encoder[-1].bias)
            
        if self.mask_encoder is not None:
            for module in self.mask_encoder:
                if isinstance(module, nn.Conv2d):
                    nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                    if module.bias is not None:
                        nn.init.zeros_(module.bias)
                elif isinstance(module, nn.GroupNorm):
                    nn.init.ones_(module.weight)
                    nn.init.zeros_(module.bias)
            # Zero-init the last layer
            nn.init.zeros_(self.mask_encoder[-1].weight)
            nn.init.zeros_(self.mask_encoder[-1].bias)
    
    def reinitialize_new_modules(self):
        """
        Reinitialize modules that were added after the base model.
        Call this after loading a pretrained model to fix uninitialized weights.
        """
        print("Reinitializing new modules (mask_predictor, mask_encoder, spatial_ref_encoder)...")
        
        # Reinitialize mask_predictor
        if self.mask_predictor is not None:
            self.mask_predictor._init_weights()
            print("  - mask_predictor reinitialized")
        
        # Reinitialize conditioning layers
        self._init_conditioning_layers()
        print("  - conditioning layers reinitialized")
        
        # Reinitialize operation embedding
        if self.operation_embedding is not None:
            nn.init.normal_(self.operation_embedding.weight, mean=0.0, std=0.02)
            print("  - operation_embedding reinitialized")
        
        # Reinitialize scalar weights
        if self.spatial_weight is not None:
            nn.init.zeros_(self.spatial_weight)
            print("  - spatial_weight reinitialized to 0")
        if self.mask_weight is not None:
            nn.init.zeros_(self.mask_weight)
            print("  - mask_weight reinitialized to 0")
        
        print("Reinitialization complete!")

    @property
    def mask_generator(self) -> EditMaskGenerator:
        """Lazy init mask generator (training only)."""
        if not self._mask_generator_initialized:
            enabled = getattr(self.config, 'mask_generator_enabled', True) and self.training
            if enabled:
                # SIMPLIFIED: Only Qwen3 + SAM3 needed now!
                self._mask_generator = EditMaskGenerator(
                    qwen_model=getattr(self.config, 'qwen_model', "Qwen/Qwen3-1.7B"),
                    device=str(self.device),
                    enabled=True
                )
            else:
                self._mask_generator = EditMaskGenerator(enabled=False)
            self._mask_generator_initialized = True
        return self._mask_generator

    def get_model(self):
        return self.model

    def mask_drop(self, latents: torch.Tensor, drop_prob: float = 0.1) -> torch.Tensor:
        if drop_prob <= 0 or not self.training:
            return latents
        mask = torch.bernoulli(torch.full((latents.shape[0],), drop_prob, device=latents.device, dtype=latents.dtype))
        while len(mask.shape) < len(latents.shape):
            mask = mask.unsqueeze(-1)
        return latents * (1 - mask)

    def get_operation_index(self, operation: str) -> int:
        if self.operation_types is None:
            return 0
        return self.operation_types.index(
            operation) if operation in self.operation_types else self.operation_types.index("other")

    def _normalize_mask(self, mask, H, W, device):
        """
        Always return a single mask: [1, H, W]
        """
        if mask is None:
            return torch.zeros(1, H, W, device=device)

        # Convert numpy → torch if needed
        if not isinstance(mask, torch.Tensor):
            mask = torch.from_numpy(mask)

        mask = mask.to(device)

        # Remove batch dim if present
        if mask.dim() == 4:      # [N, 1, H, W]
            mask = mask[:, 0]    # [N, H, W]
            # Reduction: union of all objects
            mask = mask.max(dim=0, keepdim=True)[0]

        elif mask.dim() == 3:    # [1, H, W]
            pass

        elif mask.dim() == 2:    # [H, W]
            mask = mask.unsqueeze(0)

        else:
            raise ValueError(f"Unexpected mask shape: {mask.shape}")

        return mask

    def _generate_masks_on_fly(self, und_images: torch.Tensor, instructions: List[str]) -> Tuple[
        torch.Tensor, List[str]]:
        """Generate GT masks using Qwen3 + Grounded SAM-2 (training only)."""
        masks, operations = [], []
        B, _, H, W = und_images.shape
        for i in range(und_images.shape[0]):
            try:
                mask, parsed = self.mask_generator.generate(und_images[i], instructions[i], return_parsed=True)
                mask = self._normalize_mask(mask, H=H, W=W, device=und_images.device)
                masks.append(mask)
                operations.append(parsed.get("operation", "other"))
            except Exception as e:
                print(f"Mask generation failed: {e}")
                masks.append(torch.zeros(1, H, W, device=und_images.device))
                operations.append("other")
        return torch.stack(masks).to(und_images.device), operations

    # ============================================================
    # TRAINING FORWARD
    # ============================================================
    def forward(
            self,
            input_ids: torch.LongTensor = None,
            attention_mask: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.LongTensor] = None,
            past_key_values: Optional[List[torch.FloatTensor]] = None,
            inputs_embeds: Optional[torch.FloatTensor] = None,
            labels: Optional[torch.LongTensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            gen_image: Optional[torch.FloatTensor] = None,
            und_image: Optional[torch.FloatTensor] = None,
            edit_mask: Optional[torch.FloatTensor] = None,
            operations: Optional[List[str]] = None,
            instructions: Optional[List[str]] = None,
            categories: Optional[List[str]] = None,
            return_dict: Optional[bool] = None,
            cache_position: Optional[torch.LongTensor] = None
    ) -> Union[Tuple, CausalLMOutputWithPast]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = True
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if inputs_embeds is None:
            (input_ids, position_ids, attention_mask, past_key_values,
             inputs_embeds, labels, latents) = self.prepare_inputs_labels_for_multimodal(
                input_ids, position_ids, attention_mask, past_key_values, labels, gen_image, und_image)

        # LLM Forward
        output = super().forward(
            input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
            past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels,
            use_cache=use_cache, output_attentions=output_attentions,
            output_hidden_states=output_hidden_states, return_dict=return_dict
        )

        ce_loss = output.loss
        hidden_states = output.hidden_states
        logits = output.logits
        img_hidden_states = hidden_states

        assert latents is not None

        # ============================================================
        # Generate GT Masks (Training Only)
        # ============================================================
        if edit_mask is None and instructions is not None and self.training:
            if getattr(self.config, 'mask_generator_enabled', True):
                edit_mask, operations = self._generate_masks_on_fly(und_image, instructions)

        # ============================================================
        # Predict Mask from LLM Hidden States (for inference capability)
        # ============================================================
        mask_pred_loss = torch.tensor(0.0, device=latents.device)
        predicted_mask = None
        mask_logits = None
        gt_mask_resized = None

        if self.mask_predictor is not None:
            # Get last layer hidden states
            last_hidden = hidden_states[-1]  # [B, seq_len, hidden_size]

            # Get mask logits (for stable BCE loss computation)
            mask_logits = self.mask_predictor(last_hidden, return_logits=True)  # [B, 1, H, W]

            # Resize to latent size
            mask_logits = F.interpolate(
                mask_logits.float(),
                size=(latents.shape[2], latents.shape[3]),
                mode='bilinear',
                align_corners=False
            )

            # Get probabilities for conditioning
            predicted_mask = torch.sigmoid(mask_logits)

            # Supervision loss (train predictor to match GT mask)
            if edit_mask is not None and self.training:
                gt_mask_resized = F.interpolate(
                    edit_mask.float().to(latents.device),
                    size=(latents.shape[2], latents.shape[3]),
                    mode='nearest'
                )

                # Check for NaN before loss computation
                if not torch.isnan(mask_logits).any() and not torch.isnan(gt_mask_resized).any():
                    # Standard BCE loss
                    mask_pred_loss = F.binary_cross_entropy_with_logits(
                        mask_logits,
                        gt_mask_resized,
                        reduction='mean'
                    )
                else:
                    print("WARNING: NaN in mask_logits or gt_mask, skipping mask_pred_loss")
                    mask_pred_loss = torch.tensor(0.0, device=latents.device)

        # ============================================================
        # Diffusion Setup
        # ============================================================
        noise = torch.randn_like(latents)
        weighting_scheme = "uniform"
        u = compute_density_for_timestep_sampling(
            weighting_scheme=weighting_scheme, batch_size=latents.shape[0],
            logit_mean=0.0, logit_std=1.0, mode_scale=1.29
        )
        indices = (u * self.get_model().noise_scheduler.config.num_train_timesteps).long()
        timesteps = self.get_model().noise_scheduler.timesteps[indices].to(device=latents.device)
        sigmas = self.get_sigmas(timesteps, latents.device, n_dim=latents.ndim, dtype=latents.dtype)

        # ============================================================
        # Spatial Conditioning
        # ============================================================
        if self.spatial_ref_encoder is not None:
            vae = self.get_model().get_sana_vae()
            ref_latents = vae.encode(und_image.to(vae.device)).latent * vae.config.scaling_factor
            ref_latents = ref_latents.to(latents.device)
            spatial_cond = self.spatial_ref_encoder(ref_latents)
            spatial_cond = self.mask_drop(spatial_cond, getattr(self.config, 'spatial_drop_prob', 0.1))
        else:
            spatial_cond = 0

        # ============================================================
        # Mask Conditioning (use GT mask during training)
        # ============================================================
        if self.mask_encoder is not None and edit_mask is not None:
            mask_latent = F.interpolate(
                edit_mask.float().to(latents.device),
                size=(latents.shape[2], latents.shape[3]),
                mode='nearest'
            )
            mask_latent = mask_latent.clamp(0.0, 1.0)
            
            # Do mask encoding in float32 to avoid BF16 issues
            mask_cond = mask_latent
            for layer in self.mask_encoder:
                if isinstance(layer, nn.Conv2d):
                    mask_cond = F.conv2d(mask_cond, layer.weight.float(), 
                                         layer.bias.float() if layer.bias is not None else None,
                                         layer.stride, layer.padding)
                elif isinstance(layer, nn.GroupNorm):
                    mask_cond = F.group_norm(mask_cond, layer.num_groups, 
                                             layer.weight.float(), layer.bias.float(), layer.eps)
                else:
                    mask_cond = layer(mask_cond)
            
            # Convert to model dtype and apply dropout
            mask_cond = mask_cond.to(latents.dtype)
            mask_cond = self.mask_drop(mask_cond, getattr(self.config, 'mask_drop_prob', 0.1))
        else:
            mask_cond = 0
            mask_latent = None

        # ============================================================
        # Operation Embedding
        # ============================================================
        if self.operation_embedding is not None and operations is not None:
            op_indices = torch.tensor([self.get_operation_index(op) for op in operations], device=latents.device)
            op_embed = self.operation_embedding(op_indices)[:, :, None, None]
            op_cond = op_embed * mask_latent if mask_latent is not None else op_embed.expand(-1, -1, latents.shape[2],
                                                                                             latents.shape[3])
        else:
            op_cond = 0

        # ============================================================
        # Combine Conditioning
        # ============================================================
        noisy_latents = (1.0 - sigmas) * latents + sigmas * noise
        combined_input = noisy_latents

        if self.mask_weight is not None and isinstance(mask_cond, torch.Tensor):
            combined_input = combined_input + self.mask_weight * mask_cond

        # ============================================================
        # DiT Forward
        # ============================================================
        fused_features = self.get_model().diffusion_connector(img_hidden_states)
        
        diffusion_pred = self.get_model().dit(
            hidden_states=combined_input, timestep=timesteps,
            encoder_hidden_states=fused_features, encoder_attention_mask=attention_mask
        ).sample

        target = latents - noise

        weighting = compute_loss_weighting_for_sd3(weighting_scheme=weighting_scheme, sigmas=sigmas)
        diff_loss = torch.mean(
            (weighting.float() * (diffusion_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), 1
        ).mean()

        # ============================================================
        # Total Loss
        # ============================================================
        mask_pred_weight = getattr(self.config, 'mask_predictor_loss_weight', 0.5)
        total_loss = diff_loss + 0.2 * ce_loss + mask_pred_weight * mask_pred_loss

        # Logging
        if self.training:
            print(f"Loss - diff: {diff_loss.item():.4f}, ce: {ce_loss.item():.4f}, mask_pred: {mask_pred_loss.item() if isinstance(mask_pred_loss, torch.Tensor) else 0:.4f}")

        return CausalLMOutputWithPast(
            loss=total_loss, logits=logits, past_key_values=output.past_key_values,
            hidden_states=output.hidden_states, attentions=output.attentions
        )

    # ============================================================
    # INFERENCE: Lightweight - No Qwen3/SAM-2 needed!
    # ============================================================
    @torch.no_grad()
    def generate_edited_image(
            self,
            und_image: torch.Tensor,
            input_ids: torch.Tensor,
            attention_mask: torch.Tensor,
            num_inference_steps: int = 50,
            guidance_scale: float = 7.5,
            spatial_guidance_scale: float = 1.0,
            mask_guidance_scale: float = 1.0,
            generator: Optional[torch.Generator] = None,
    ) -> torch.Tensor:
        """
        Lightweight inference - uses learned mask predictor instead of SAM-2.

        Args:
            und_image: Input image tensor [B, C, H, W]
            input_ids: Tokenized prompt [B, seq_len]
            attention_mask: Attention mask [B, seq_len]
            num_inference_steps: Denoising steps
            guidance_scale: CFG scale for text
            spatial_guidance_scale: Scale for spatial conditioning
            mask_guidance_scale: Scale for predicted mask conditioning
            generator: Random generator for reproducibility

        Returns:
            Edited image latents [B, C, H, W]
        """

        device = und_image.device
        dtype = und_image.dtype
        batch_size = und_image.shape[0]

        # ============================================================
        # 1. Get LLM Hidden States
        # ============================================================
        (input_ids_mm, position_ids, attention_mask_mm, _,
         inputs_embeds, _, _) = self.prepare_inputs_labels_for_multimodal(
            input_ids, None, attention_mask, None, None, None, und_image
        )

        output = Qwen2ForCausalLM.forward(
            self,
            input_ids=input_ids_mm,
            attention_mask=attention_mask_mm,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            output_hidden_states=True,
            return_dict=True
        )

        hidden_states = output.hidden_states
        img_hidden_states = hidden_states

        # ============================================================
        # 2. Predict Edit Mask (NO SAM-2 needed!)
        # ============================================================
        if self.mask_predictor is not None:
            last_hidden = hidden_states[-1]
            predicted_mask = self.mask_predictor(last_hidden)  # [B, 1, H, W]
        else:
            predicted_mask = None

        # ============================================================
        # 3. Encode Reference Image
        # ============================================================
        vae = self.get_model().get_sana_vae()
        ref_latents = vae.encode(und_image.to(vae.device)).latent * vae.config.scaling_factor
        ref_latents = ref_latents.to(device)

        latent_h, latent_w = ref_latents.shape[2], ref_latents.shape[3]
        latent_channels = ref_latents.shape[1]

        # Resize predicted mask to latent size
        if predicted_mask is not None:
            predicted_mask = F.interpolate(
                predicted_mask, size=(latent_h, latent_w), mode='bilinear', align_corners=False
            )

        # ============================================================
        # 4. Prepare Conditioning
        # ============================================================
        # Spatial conditioning
        if self.spatial_ref_encoder is not None:
            spatial_cond = self.spatial_ref_encoder(ref_latents)
        else:
            spatial_cond = torch.zeros_like(ref_latents)

        # Mask conditioning
        if self.mask_encoder is not None and predicted_mask is not None:
            mask_cond = self.mask_encoder(predicted_mask.to(dtype=self.mask_encoder[0].weight.dtype))
        else:
            mask_cond = torch.zeros_like(ref_latents)

        # Semantic conditioning from LLM
        fused_features = self.get_model().diffusion_connector(img_hidden_states)

        # ============================================================
        # 5. Prepare for CFG
        # ============================================================
        if guidance_scale > 1.0:
            # Unconditional: zero out conditioning
            spatial_cond_uncond = torch.zeros_like(spatial_cond)
            mask_cond_uncond = torch.zeros_like(mask_cond)
            fused_features_uncond = torch.zeros_like(fused_features)

            # Stack [uncond, cond]
            spatial_cond_cfg = torch.cat([spatial_cond_uncond, spatial_cond])
            mask_cond_cfg = torch.cat([mask_cond_uncond, mask_cond])
            fused_features_cfg = torch.cat([fused_features_uncond, fused_features])
        else:
            spatial_cond_cfg = spatial_cond
            mask_cond_cfg = mask_cond
            fused_features_cfg = fused_features

        # ============================================================
        # 6. Initialize Latents
        # ============================================================
        latents = randn_tensor(
            (batch_size, latent_channels, latent_h, latent_w),
            generator=generator, device=device, dtype=dtype
        )

        # ============================================================
        # 7. Setup Scheduler
        # ============================================================
        scheduler = self.get_model().noise_scheduler
        scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = scheduler.timesteps

        # ============================================================
        # 8. Denoising Loop
        # ============================================================
        for t in timesteps:
            # Expand for CFG
            if guidance_scale > 1.0:
                latent_model_input = torch.cat([latents] * 2)
                t_input = torch.cat([t.unsqueeze(0)] * 2 * batch_size)
            else:
                latent_model_input = latents
                t_input = t.unsqueeze(0).expand(batch_size)

            # Add conditioning
            combined_input = latent_model_input
            if self.spatial_weight is not None:
                combined_input = combined_input + spatial_guidance_scale * self.spatial_weight * spatial_cond_cfg
            if self.mask_weight is not None:
                combined_input = combined_input + mask_guidance_scale * self.mask_weight * mask_cond_cfg

            # DiT forward
            noise_pred = self.get_model().dit(
                hidden_states=combined_input,
                timestep=t_input,
                encoder_hidden_states=fused_features_cfg,
            ).sample

            # CFG
            if guidance_scale > 1.0:
                noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)

            # Scheduler step
            latents = scheduler.step(noise_pred, t, latents).prev_sample

        # ============================================================
        # 9. Decode Latents
        # ============================================================
        latents = latents / vae.config.scaling_factor
        image = vae.decode(latents.to(vae.device)).sample

        return image


# ============================================================
# Register Model
# ============================================================

AutoConfig.register("llava_qwen2", blip3oFastConfig)
AutoModelForCausalLM.register(blip3oFastConfig, blip3oFastForCausalLM)