code/src/models/VLANeXt.py

import sys
import os

from PIL import Image
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    AutoProcessor, AutoTokenizer,
    SiglipVisionModel, SiglipImageProcessor, LlamaForCausalLM, 
    PaliGemmaForConditionalGeneration, 
    Qwen3VLForConditionalGeneration
)
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler

from .policies import (
    ActionDiffusionTransformerMetaquery, ActionDiffusionTransformerMoE,
    ActionRegressionTransformerMetaquery, ActionRegressionTransformerMoE,
    ActionClassificationTransformerMetaquery, ActionClassificationTransformerMoE, ActionVQVAE
)
from .generator import ImageGeneratorTransformer
from .encoder import ActionTransformerProjector
from .connector import ConnectorTransformer

try:
    from .Emu3_5_VisionTokenizer.modeling_emu3p5visionvq import Emu3p5VisionVQModel
except ImportError:
    # Fallback for directory with dot in name (Emu3.5_VisionTokenizer) which is not a valid package name
    sys.path.append(os.path.join(os.path.dirname(__file__), "Emu3.5_VisionTokenizer"))
    from modeling_emu3p5visionvq import Emu3p5VisionVQModel



class LlamaProcessorWrapper:
    def __init__(self, tokenizer, image_processor):
        self.tokenizer = tokenizer
        self.image_processor = image_processor

class VLANeXt(nn.Module):
    def __init__(
        self, 
        lmm_path="Qwen/Qwen3-VL-2B-Instruct",
        vision_encoder_path="google/siglip2-base-patch16-256",
        action_dim=7,
        num_actions=1,
        num_queries=16,
        num_history=0,
        loss_type="diffusion", # Options: "diffusion", "regression", "classification"
        future_image_loss_weight=0.0,
        num_train_timesteps=1000,
        num_inference_timesteps=10,
        scheduler_type="ddim", # Options: "ddim", "flow_match"
        condition_type="loose", # Options: "loose", "tight", "soft"
        policy_hidden_size=1024,
        policy_depth=24,
        policy_num_heads=16,
        policy_mlp_ratio=4.0,
        policy_mixer_type="attention",
        policy_mix_every_n=4,
        policy_ttt_causal=False,
        policy_ttt_chunk_size=64,
        use_proprio_input_vlm=True,
        use_action_input_policy=False,
        use_transformer_proprio_projector=True,
        projector_depth=2,
        projector_num_heads=4,
        use_transformer_connector=True,
        connector_depth=2,
        connector_num_heads=4,
        backbone_mode="finetune", # Options: "frozen", "finetune"
        gradient_checkpointing=True,
        num_bins=256,
        action_vqvae=None,

        generator_hidden_size=768,
        generator_depth=12,
        generator_num_heads=12,
        generator_mlp_ratio=4.0,
        generator_mixer_type="attention",
        generator_mix_every_n=4,
        generator_ttt_chunk_size=16,
        generator_fallback_mixer="attention",
        generator_swa_window_size=64,
        generator_ttt_use_cuda_kernel=False,
        attn_implementation="flash_attention_2",
        dct_loss_weight=0.1,
        dct_low_freq_weight=1.0,
        dct_high_freq_weight=3.0,
        dct_freq_split=0.5,
        dct_similarity_type="mse",  # Options: "mse", "mae", "cosine"
    ):
        super().__init__()
        
        print(f"Initializing VLM {lmm_path} with attn_implementation: {attn_implementation}")
        if "paligemma" in lmm_path.lower():
            self.model_family = "paligemma"
            self.lmm = PaliGemmaForConditionalGeneration.from_pretrained(
                lmm_path, dtype=torch.bfloat16, _attn_implementation=attn_implementation
            )
            self.processor = AutoProcessor.from_pretrained(lmm_path, trust_remote_code=True)
            if hasattr(self.lmm.config, "text_config"):
                self.hidden_size = self.lmm.config.text_config.hidden_size
            else:
                self.hidden_size = self.lmm.config.hidden_size
        elif "llama" in lmm_path.lower():
            self.model_family = "llama"
            self.lmm = LlamaForCausalLM.from_pretrained(
                lmm_path, dtype=torch.bfloat16, attn_implementation=attn_implementation
            )
            self.vision_encoder = SiglipVisionModel.from_pretrained(
                vision_encoder_path, dtype=torch.bfloat16, attn_implementation=attn_implementation
            )
            tokenizer = AutoTokenizer.from_pretrained(lmm_path)
            if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
            image_processor = SiglipImageProcessor.from_pretrained(vision_encoder_path)
            self.processor = LlamaProcessorWrapper(tokenizer, image_processor)
            self.hidden_size = self.lmm.config.hidden_size
            self.vision_projector = nn.Sequential(
                nn.Linear(self.vision_encoder.config.hidden_size, self.hidden_size),
                nn.LayerNorm(self.hidden_size),
                nn.SiLU(), 
                nn.Linear(self.hidden_size, self.hidden_size),
                nn.LayerNorm(self.hidden_size),
                nn.SiLU(), 
                nn.Linear(self.hidden_size, self.hidden_size)
            )
        elif "qwen" in lmm_path.lower():
            self.model_family = "qwen"
            self.lmm = Qwen3VLForConditionalGeneration.from_pretrained(
                lmm_path, dtype=torch.bfloat16, _attn_implementation=attn_implementation
            )
            self.processor = AutoProcessor.from_pretrained(lmm_path, trust_remote_code=True)
            if hasattr(self.lmm.config, "text_config"):
                self.hidden_size = self.lmm.config.text_config.hidden_size
            else:
                self.hidden_size = self.lmm.config.hidden_size
        
        if backbone_mode == "frozen":
            self.lmm.requires_grad_(False)
            if self.model_family == "llama":
                self.vision_encoder.requires_grad_(False)
        elif backbone_mode == "finetune":
            self.lmm.requires_grad_(True)
            if self.model_family == "llama":
                self.vision_encoder.requires_grad_(True)
        else:
            raise ValueError(f"Unknown backbone_mode: {backbone_mode}")

        if gradient_checkpointing:
            model_to_configure = self.lmm
            if hasattr(model_to_configure, "gradient_checkpointing_enable"):
                model_to_configure.gradient_checkpointing_enable()
            if hasattr(self.lmm, "enable_input_require_grads"):
                self.lmm.enable_input_require_grads()
            config = self.lmm.config
            if hasattr(config, "use_cache"):
                config.use_cache = False
            if self.model_family == "llama":
                 if hasattr(self.vision_encoder, "gradient_checkpointing_enable"):
                    self.vision_encoder.gradient_checkpointing_enable()

        self.num_queries = num_queries
        self.loss_type = loss_type
        self.scheduler_type = scheduler_type
        self.num_train_timesteps = num_train_timesteps
        self.num_inference_timesteps = num_inference_timesteps
        self.action_dim = action_dim
        self.num_actions = num_actions
        self.num_history = num_history
        self.num_bins = num_bins
        self.condition_type = condition_type
        self.use_proprio_input_vlm = use_proprio_input_vlm
        self.use_action_input_policy = use_action_input_policy
        self.future_image_loss_weight = future_image_loss_weight
        self.enable_future_image_loss = (future_image_loss_weight > 0)
        self.dct_loss_weight = dct_loss_weight
        self.dct_low_freq_weight = dct_low_freq_weight
        self.dct_high_freq_weight = dct_high_freq_weight
        self.dct_freq_split = dct_freq_split
        self.dct_similarity_type = dct_similarity_type
        
        self.action_vqvae_config = action_vqvae
        if self.action_vqvae_config.get('enabled', False):
            self.action_vqvae = ActionVQVAE(
                action_dim=action_dim,
                latent_codes_per_step=3, 
                codebook_size=self.action_vqvae_config.get('codebook_size', 1024),
                hidden_size=self.action_vqvae_config.get('hidden_size', 256),
                depth=self.action_vqvae_config.get('depth', 2),
                num_heads=self.action_vqvae_config.get('num_heads', 4)
            )
        else:
            self.action_vqvae = None


        if self.enable_future_image_loss:
            print("Initializing Future Image Generator Components...")
            self.vq_model = Emu3p5VisionVQModel.from_pretrained("/mnt/afs-h200/yuyangcheng/models/Emu3.5-VisionTokenizer", trust_remote_code=True)
            self.vq_model.requires_grad_(False)
            self.vq_codebook_size = self.vq_model.config.codebook_size
            
            self.generator = ImageGeneratorTransformer(
                vocab_size=self.vq_codebook_size,
                vlm_hidden_size=self.hidden_size,
                hidden_size=generator_hidden_size,
                depth=generator_depth,
                num_heads=generator_num_heads,
                mlp_ratio=generator_mlp_ratio,
                mixer_type=generator_mixer_type,
                mix_every_n=generator_mix_every_n,
                ttt_chunk_size=generator_ttt_chunk_size,
                fallback_mixer=generator_fallback_mixer,
                swa_window_size=generator_swa_window_size,
                ttt_use_cuda_kernel=generator_ttt_use_cuda_kernel,
            )
        else:
            self.vq_model = None
            self.generator = None

        # When True, predict_action decodes the 256 image tokens with the
        # generator's O(n) incremental path (generate_incremental) instead of the
        # O(n^2) full-recompute AR loop. Default False -> byte-identical to the
        # original loop (zero eval regression). Toggled by eval/bench scripts.
        self.use_incremental_gen = False

        if self.use_proprio_input_vlm:
            projector_input_dim = action_dim
            if use_transformer_proprio_projector:
                self.action_projector = ActionTransformerProjector(
                    action_dim=projector_input_dim,
                    hidden_size=self.hidden_size,
                    depth=projector_depth,
                    num_heads=projector_num_heads
                )
            else:
                self.action_projector = nn.Linear(projector_input_dim, self.hidden_size)
        else:
            self.action_projector = None
        
        self.meta_queries = nn.Parameter(
            torch.randn(num_queries, self.hidden_size)
        )
        if self.condition_type == "loose":
            if use_transformer_connector:
                self.connector = ConnectorTransformer(
                    input_dim=self.hidden_size,
                    output_dim=self.hidden_size,
                    depth=connector_depth,
                    num_heads=connector_num_heads
                )
            else:
                self.connector = nn.Sequential(
                    nn.Linear(self.hidden_size, self.hidden_size),
                    nn.SiLU(),
                    nn.Linear(self.hidden_size, self.hidden_size) # Project to diffusion cond dim
                )
        else:
            self.connector = None

        gen_hidden_dim = generator_hidden_size if self.enable_future_image_loss else None
        if loss_type == "regression":
            if condition_type in ["tight", "soft"]:
                self.action_head = ActionRegressionTransformerMoE(
                    action_dim=action_dim,
                    vlm_hidden_size=self.hidden_size,
                    num_actions=num_actions,
                    hidden_size=policy_hidden_size,
                    depth=policy_depth,
                    num_heads=policy_num_heads,
                    mlp_ratio=policy_mlp_ratio,
                    gen_hidden_size=gen_hidden_dim,
                    mixer_type=policy_mixer_type,
                    mix_every_n=policy_mix_every_n,
                    ttt_causal=policy_ttt_causal,
                    ttt_chunk_size=policy_ttt_chunk_size,
                )
            elif condition_type == "loose":
                self.action_head = ActionRegressionTransformerMetaquery(
                    action_dim=action_dim,
                    condition_dim=self.hidden_size,
                    num_actions=num_actions,
                    hidden_size=policy_hidden_size,
                    depth=policy_depth,
                    num_heads=policy_num_heads,
                    mlp_ratio=policy_mlp_ratio
                )
            else:
                raise ValueError(f"Unknown condition type for regression: {condition_type}")
            self.noise_scheduler = None
        elif loss_type == "classification":
            is_vqvae = (self.action_vqvae is not None)

            if condition_type == "loose":
                if is_vqvae:
                    self.action_head = ActionClassificationTransformerMetaquery(
                        action_dim=action_dim,
                        condition_dim=self.hidden_size,
                        num_actions=num_actions,
                        hidden_size=policy_hidden_size,
                        depth=policy_depth,
                        num_heads=policy_num_heads,
                        mlp_ratio=policy_mlp_ratio,
                        vqvae_mode=True,
                        vq_codebook_size=self.action_vqvae.codebook_size,
                        vq_latent_codes=self.action_vqvae.latent_codes
                    )
                else:
                    self.action_head = ActionClassificationTransformerMetaquery(
                        action_dim=action_dim,
                        condition_dim=self.hidden_size,
                        num_actions=num_actions,
                        num_bins=num_bins,
                        hidden_size=policy_hidden_size,
                        depth=policy_depth,
                        num_heads=policy_num_heads,
                        mlp_ratio=policy_mlp_ratio,
                        vqvae_mode=False
                    )
            elif condition_type in ["tight", "soft"]:
                if is_vqvae:
                    self.action_head = ActionClassificationTransformerMoE(
                        action_dim=action_dim,
                        vlm_hidden_size=self.hidden_size,
                        num_actions=num_actions,
                        hidden_size=policy_hidden_size,
                        depth=policy_depth,
                        num_heads=policy_num_heads,
                        mlp_ratio=policy_mlp_ratio,
                        vqvae_mode=True,
                        vq_codebook_size=self.action_vqvae.codebook_size,
                        vq_latent_codes=self.action_vqvae.latent_codes,
                        gen_hidden_size=gen_hidden_dim,
                        mixer_type=policy_mixer_type,
                        mix_every_n=policy_mix_every_n,
                        ttt_causal=policy_ttt_causal,
                        ttt_chunk_size=policy_ttt_chunk_size,
                    )
                else:
                    self.action_head = ActionClassificationTransformerMoE(
                        action_dim=action_dim,
                        vlm_hidden_size=self.hidden_size,
                        num_actions=num_actions,
                        num_bins=num_bins,
                        hidden_size=policy_hidden_size,
                        depth=policy_depth,
                        num_heads=policy_num_heads,
                        mlp_ratio=policy_mlp_ratio,
                        vqvae_mode=False,
                        gen_hidden_size=gen_hidden_dim,
                        mixer_type=policy_mixer_type,
                        mix_every_n=policy_mix_every_n,
                        ttt_causal=policy_ttt_causal,
                        ttt_chunk_size=policy_ttt_chunk_size,
                    )
            else:
                raise NotImplementedError(f"Classification policy does not support {condition_type}.")
            self.noise_scheduler = None
        elif loss_type == "diffusion":
            if condition_type in ["tight", "soft"]:
                self.action_head = ActionDiffusionTransformerMoE(
                    action_dim=action_dim,
                    vlm_hidden_size=self.hidden_size,
                    hidden_size=policy_hidden_size,
                    depth=policy_depth,
                    num_heads=policy_num_heads,
                    mlp_ratio=policy_mlp_ratio,
                    gen_hidden_size=gen_hidden_dim,
                    mixer_type=policy_mixer_type,
                    mix_every_n=policy_mix_every_n,
                    ttt_causal=policy_ttt_causal,
                    ttt_chunk_size=policy_ttt_chunk_size,
                )
            elif condition_type == "loose":
                self.action_head = ActionDiffusionTransformerMetaquery(
                    action_dim=action_dim,
                    condition_dim=self.hidden_size,
                    hidden_size=policy_hidden_size,
                    depth=policy_depth,
                    num_heads=policy_num_heads,
                    mlp_ratio=policy_mlp_ratio
                )
            else:
                raise ValueError(f"Unknown condition type for diffusion: {condition_type}")
        else:
            raise ValueError(f"Unknown loss type: {loss_type}")

        if loss_type == "diffusion":
            if scheduler_type == "ddim":
                self.noise_scheduler = DDIMScheduler(
                    num_train_timesteps=num_train_timesteps,
                    clip_sample=False,
                    prediction_type="epsilon"
                )
            elif scheduler_type == "flow_match": 
                self.noise_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=num_train_timesteps)
            else:
                raise ValueError(f"Unknown scheduler type: {scheduler_type}")

    def forward_action_vqvae_pretrain(self, actions):
        if self.action_vqvae is None:
            raise RuntimeError("Action VQ-VAE not initialized.")
            
        actions = actions.to(dtype=self.action_vqvae.in_proj.weight.dtype)
        loss = self.action_vqvae(actions)
        return loss

    def get_vlm_condition(self, input_ids, attention_mask, proprioception=None, proprio_attention_mask=None, pixel_values=None, pixel_values_videos=None, image_grid_thw=None, video_grid_thw=None, token_type_ids=None):
        if self.model_family == "paligemma":
            return self._get_vlm_condition_paligemma(input_ids, attention_mask, proprioception, proprio_attention_mask, pixel_values, token_type_ids=token_type_ids)
        elif self.model_family == "llama":
            return self._get_vlm_condition_llama(input_ids, attention_mask, pixel_values, proprioception, proprio_attention_mask)
        elif self.model_family == "qwen":
            return self._get_vlm_condition_qwen(input_ids, attention_mask, proprioception, proprio_attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw)

    def _get_vlm_condition_qwen(self, input_ids, attention_mask, proprioception, proprio_attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw):
        B = input_ids.shape[0]
        
        backbone = self.lmm.model
        lmm_config = self.lmm.config
        pad_token_id = getattr(lmm_config, "pad_token_id", None)
        pad_token_id = pad_token_id if pad_token_id is not None else 0
        inputs_embeds = backbone.get_input_embeddings()(input_ids)
        
        if self.use_proprio_input_vlm and proprioception is not None:
            proprio_embeds = self.action_projector(proprioception.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype))
            inputs_embeds = torch.cat([proprio_embeds, inputs_embeds], dim=1)
            if attention_mask is not None:
                if proprio_attention_mask is not None:
                    proprio_mask = proprio_attention_mask.to(device=attention_mask.device, dtype=attention_mask.dtype)
                else:
                    proprio_mask = torch.ones(B, proprioception.shape[1], device=attention_mask.device, dtype=attention_mask.dtype)
                attention_mask = torch.cat([proprio_mask, attention_mask], dim=1)
            proprio_ids = torch.full((B, proprioception.shape[1]), pad_token_id, dtype=input_ids.dtype, device=input_ids.device)
            input_ids = torch.cat([proprio_ids, input_ids], dim=1)

        if self.condition_type != "tight":
            queries_embeds = self.meta_queries.unsqueeze(0).expand(B, -1, -1).to(inputs_embeds.dtype)
            inputs_embeds = torch.cat([inputs_embeds, queries_embeds], dim=1)
            if attention_mask is not None:
                queries_mask = torch.ones(B, self.num_queries, device=attention_mask.device, dtype=attention_mask.dtype)
                attention_mask = torch.cat([attention_mask, queries_mask], dim=1)
            queries_ids = torch.full((B, self.num_queries), pad_token_id, dtype=input_ids.dtype, device=input_ids.device)
            extended_input_ids = torch.cat([input_ids, queries_ids], dim=1)
        else:
            extended_input_ids = input_ids

        rope_kwargs = {
            "input_ids": extended_input_ids,
            "image_grid_thw": image_grid_thw,
            "video_grid_thw": video_grid_thw,
            "attention_mask": attention_mask
        }

        position_ids, _ = backbone.get_rope_index(**rope_kwargs)

        output_hidden_states_flag = (self.enable_future_image_loss or self.condition_type in ["tight", "soft"])
        forward_kwargs = {
            "inputs_embeds": inputs_embeds,
            "position_ids": position_ids,
            "attention_mask": attention_mask,
            "pixel_values": pixel_values,
            "pixel_values_videos": pixel_values_videos,
            "image_grid_thw": image_grid_thw,
            "video_grid_thw": video_grid_thw,
            "output_hidden_states": output_hidden_states_flag,
        }
        outputs = backbone(**forward_kwargs)
        hidden_states = outputs.hidden_states if output_hidden_states_flag else None
        connector_out = None
        if self.condition_type == "loose" and self.connector is not None:
            query_outputs = outputs.last_hidden_state[:, -self.num_queries:, :]
            connector_out = self.connector(query_outputs)
            
        return connector_out, hidden_states

    def _get_vlm_condition_llama(self, input_ids, attention_mask, pixel_values, proprioception, proprio_attention_mask):
        B = input_ids.shape[0]
        pixel_values = pixel_values.to(dtype=self.vision_encoder.dtype)
        
        vision_outputs = self.vision_encoder(pixel_values, output_hidden_states=True)
        image_feats = vision_outputs.last_hidden_state
        image_embeds = self.vision_projector(image_feats)

        if image_embeds.shape[0] != B:
            num_views = image_embeds.shape[0] // B
            image_embeds = image_embeds.view(B, num_views, -1, image_embeds.shape[-1])
            image_embeds = image_embeds.flatten(1, 2)
        
        text_embeds = self.lmm.model.embed_tokens(input_ids)

        proprio_embeds = None
        if self.use_proprio_input_vlm and proprioception is not None:
             proprio_embeds = self.action_projector(proprioception.to(device=text_embeds.device, dtype=text_embeds.dtype))

        embeds_list = [image_embeds]
        image_mask = torch.ones(B, image_embeds.shape[1], device=attention_mask.device, dtype=attention_mask.dtype)
        mask_list = [image_mask]

        if proprio_embeds is not None:
            embeds_list.append(proprio_embeds)
            if proprio_attention_mask is not None:
                mask_list.append(proprio_attention_mask.to(attention_mask.device))
            else:
                p_mask = torch.ones(B, proprio_embeds.shape[1], device=attention_mask.device, dtype=attention_mask.dtype)
                mask_list.append(p_mask)
        
        embeds_list.append(text_embeds)
        mask_list.append(attention_mask)

        if self.condition_type != "tight":
            queries_embeds = self.meta_queries.unsqueeze(0).expand(B, -1, -1).to(text_embeds.dtype)
            embeds_list.append(queries_embeds)
            queries_mask = torch.ones(B, self.num_queries, device=attention_mask.device, dtype=attention_mask.dtype)
            mask_list.append(queries_mask)

        inputs_embeds = torch.cat(embeds_list, dim=1)
        combined_attention_mask = torch.cat(mask_list, dim=1)

        output_hidden_states_flag = (self.enable_future_image_loss or self.condition_type in ["tight", "soft"])
        outputs = self.lmm.model(
            inputs_embeds=inputs_embeds,
            attention_mask=combined_attention_mask,
            output_hidden_states=output_hidden_states_flag
        )
        hidden_states = outputs.hidden_states if output_hidden_states_flag else None
        connector_out = None
        if self.condition_type == "loose" and self.connector is not None:
            query_outputs = outputs.last_hidden_state[:, -self.num_queries:, :]
            connector_out = self.connector(query_outputs)

        return connector_out, hidden_states

    def _get_vlm_condition_paligemma(self, input_ids, attention_mask, proprioception, proprio_attention_mask, pixel_values, token_type_ids=None):
        from transformers.models.paligemma.modeling_paligemma import create_causal_mask_mapping

        B = input_ids.shape[0]

        backbone = self.lmm.model

        inputs_embeds = backbone.get_input_embeddings()(input_ids)

        if pixel_values is not None:
            image_outputs = backbone.get_image_features(pixel_values)
            image_features = image_outputs.pooler_output
            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
            special_image_mask = backbone.get_placeholder_mask(input_ids, inputs_embeds, image_features)
            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)

        if self.use_proprio_input_vlm and proprioception is not None:
            proprio_embeds = self.action_projector(proprioception.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype))
            inputs_embeds = torch.cat([proprio_embeds, inputs_embeds], dim=1)
            if attention_mask is not None:
                if proprio_attention_mask is not None:
                    proprio_mask = proprio_attention_mask.to(device=attention_mask.device, dtype=attention_mask.dtype)
                else:
                    proprio_mask = torch.ones(B, proprioception.shape[1], device=attention_mask.device, dtype=attention_mask.dtype)
                attention_mask = torch.cat([proprio_mask, attention_mask], dim=1)
            # Proprio tokens are prefix context — token_type_ids=0 (bidirectional)
            if token_type_ids is not None:
                proprio_type_ids = torch.zeros(B, proprioception.shape[1], device=token_type_ids.device, dtype=token_type_ids.dtype)
                token_type_ids = torch.cat([proprio_type_ids, token_type_ids], dim=1)

        if self.condition_type != "tight":
            queries_embeds = self.meta_queries.unsqueeze(0).expand(B, -1, -1).to(inputs_embeds.dtype)
            inputs_embeds = torch.cat([inputs_embeds, queries_embeds], dim=1)
            if attention_mask is not None:
                queries_mask = torch.ones(B, self.num_queries, device=attention_mask.device, dtype=attention_mask.dtype)
                attention_mask = torch.cat([attention_mask, queries_mask], dim=1)
            # Query tokens are suffix — token_type_ids=1 (causal)
            if token_type_ids is not None:
                queries_type_ids = torch.ones(B, self.num_queries, device=token_type_ids.device, dtype=token_type_ids.dtype)
                token_type_ids = torch.cat([token_type_ids, queries_type_ids], dim=1)

        # Build the proper PaliGemma causal mask with bidirectional attention on prefix/image tokens
        cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)
        position_ids = cache_position.unsqueeze(0) + 1  # PaliGemma positions are 1-indexed
        causal_mask_mapping = create_causal_mask_mapping(
            backbone.config,
            inputs_embeds,
            attention_mask,
            cache_position,
            past_key_values=None,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            pixel_values=pixel_values,
            is_training=self.training,
        )

        output_hidden_states_flag = (self.enable_future_image_loss or self.condition_type in ["tight", "soft"] )
        outputs = backbone.language_model(
            inputs_embeds=inputs_embeds,
            attention_mask=causal_mask_mapping,
            position_ids=position_ids,
            output_hidden_states=output_hidden_states_flag,
        )
        hidden_states = outputs.hidden_states if output_hidden_states_flag else None
        connector_out = None
        if self.condition_type == "loose" and self.connector is not None:
            query_outputs = outputs.last_hidden_state[:, -self.num_queries:, :]
            connector_out = self.connector(query_outputs)

        return connector_out, hidden_states

    def _compute_gen_loss_and_feats(self, future_images, vlm_hidden_states):
        with torch.no_grad():
            future_images = future_images.to(device=self.vq_model.device, dtype=self.vq_model.dtype)
            _, _, (_, _, token_ids) = self.vq_model.encode(future_images)
            B = future_images.shape[0]
            token_ids = token_ids.view(B, -1)
        
        sos_token = torch.zeros((B, 1), dtype=token_ids.dtype, device=token_ids.device)
        gen_input = torch.cat([sos_token, token_ids[:, :-1]], dim=1)
        
        gen_logits, gen_hidden_states = self.generator(gen_input, vlm_hidden_states)
        loss_img = F.cross_entropy(gen_logits.reshape(-1, self.vq_codebook_size), token_ids.reshape(-1))
        
        return loss_img, gen_hidden_states

    def _compute_dct_loss(self, pred, target):
        B, T, D = pred.shape

        if not hasattr(self, '_dct_matrix') or self._dct_matrix.shape[0] != T or self._dct_matrix.device != pred.device:
            n = torch.arange(T, device=pred.device).float()
            k = torch.arange(T, device=pred.device).float()
            dct_m = torch.cos((np.pi / T) * (n + 0.5).unsqueeze(0) * k.unsqueeze(1))
            
            dct_m[0, :] *= 1.0 / np.sqrt(T)
            dct_m[1:, :] *= np.sqrt(2.0 / T)
            
            self._dct_matrix = dct_m

        split_idx = max(1, int(T * self.dct_freq_split))
        freq_weights = torch.ones(T, device=pred.device, dtype=pred.dtype)
        freq_weights[:split_idx] = self.dct_low_freq_weight
        freq_weights[split_idx:] = self.dct_high_freq_weight
        freq_weights = freq_weights.view(1, T, 1)

        pred_perm = pred.permute(0, 2, 1)
        pred_dct = torch.matmul(pred_perm, self._dct_matrix.t())
        pred_dct = pred_dct.permute(0, 2, 1)

        target_perm = target.permute(0, 2, 1)
        target_dct = torch.matmul(target_perm, self._dct_matrix.t())
        target_dct = target_dct.permute(0, 2, 1)

        sim_type = self.dct_similarity_type
        if sim_type == "mse":
            diff = (pred_dct - target_dct) ** 2
            return (diff * freq_weights).mean()
        elif sim_type == "mae":
            diff = (pred_dct - target_dct).abs()
            return (diff * freq_weights).mean()
        elif sim_type == "cosine":
            pred_norm = torch.nn.functional.normalize(pred_dct, dim=-1)
            target_norm = torch.nn.functional.normalize(target_dct, dim=-1)
            cos_sim = (pred_norm * target_norm).sum(dim=-1, keepdim=True)
            cos_dist = 1.0 - cos_sim
            return (cos_dist * freq_weights).mean()
        else:
            raise ValueError(f"Unknown dct_similarity_type: {sim_type!r}. "
                             f"Options are: 'mse', 'mae', 'cosine'.")

    def forward(self, input_ids=None, attention_mask=None, actions=None, proprioception=None, history_actions=None, proprio_attention_mask=None, pixel_values=None, pixel_values_videos=None, image_grid_thw=None, video_grid_thw=None, future_images=None, task=None, token_type_ids=None):
        if task == "action_vqvae_pretrain":
            return self.forward_action_vqvae_pretrain(actions)
            
        if self.loss_type == "regression":
            return self._forward_regression(
                input_ids, attention_mask, actions, proprioception, history_actions, proprio_attention_mask,
                pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw, future_images, token_type_ids=token_type_ids
            )
        elif self.loss_type == "classification":
            return self._forward_classification(
                input_ids, attention_mask, actions, proprioception, history_actions, proprio_attention_mask,
                pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw, future_images, token_type_ids=token_type_ids
            )
        elif self.loss_type == "diffusion":
            return self._forward_diffusion(
                input_ids, attention_mask, actions, proprioception, history_actions, proprio_attention_mask,
                pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw, future_images, token_type_ids=token_type_ids
            )

    def _forward_classification(self, input_ids, attention_mask, actions, proprioception, history_actions, proprio_attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw, future_images=None, token_type_ids=None):
        connector_out, hidden_states = self.get_vlm_condition(
            input_ids, attention_mask, proprioception=proprioception, proprio_attention_mask=proprio_attention_mask,
            pixel_values=pixel_values, pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw,
            token_type_ids=token_type_ids
        )
        
        loss_img = 0.0
        gen_hidden_states = None
        if self.enable_future_image_loss and future_images is not None:
             loss_img, gen_hidden_states = self._compute_gen_loss_and_feats(future_images, hidden_states)

        policy_history = history_actions if self.use_action_input_policy else None
        
        if self.condition_type in ["tight", "soft"]:
            if self.enable_future_image_loss:
                pred_logits = self.action_head(hidden_states, history_actions=policy_history, gen_hidden_states=gen_hidden_states)
            else:
                pred_logits = self.action_head(hidden_states, history_actions=policy_history)
        elif self.condition_type == "loose":
            cond_input = connector_out.mean(dim=1)
            pred_logits = self.action_head(cond_input, history_actions=policy_history)
        else:
            raise ValueError(f"Unknown condition type: {self.condition_type}")
        
        if actions.ndim == 2: actions = actions.unsqueeze(1)

        pred_action_continuous = None
        loss = 0.0

        if self.action_vqvae is not None:
             with torch.no_grad():
                 self.action_vqvae.eval()
                 actions_input = actions.to(dtype=self.action_vqvae.in_proj.weight.dtype)
                 _, indices, _ = self.action_vqvae.encode(actions_input)
             
             loss = F.cross_entropy(
                 pred_logits.reshape(-1, self.action_vqvae.codebook_size), 
                 indices.reshape(-1)
             )

             if self.dct_loss_weight > 0:
                 probs = F.softmax(pred_logits, dim=-1)
                 pred_action_continuous = self.action_vqvae.decode_probs(probs)
        else:
            logits = pred_logits
            pose_logits = logits[:, :, :self.action_dim - 1, :]
            gripper_logits = logits[:, :, -1:, :2]
            
            gt_pose = torch.clamp(actions[:, :, :6], -1, 1)
            gt_pose_idx = ((gt_pose + 1) / 2 * (self.num_bins - 1)).round().long()
            
            gt_gripper = torch.clamp(actions[:, :, 6:7], -1, 1)
            gt_gripper_idx = ((gt_gripper + 1) / 2).round().long() # 0 or 1

            loss_pose = F.cross_entropy(pose_logits.reshape(-1, self.num_bins), gt_pose_idx.reshape(-1))
            loss_gripper = F.cross_entropy(gripper_logits.reshape(-1, 2), gt_gripper_idx.reshape(-1))
            
            loss = (loss_pose + loss_gripper) / 2.0

            if self.dct_loss_weight > 0:
                pose_probs = F.softmax(pose_logits, dim=-1)
                bin_centers = torch.linspace(-1, 1, self.num_bins, device=actions.device, dtype=pose_probs.dtype)
                pred_pose = torch.sum(pose_probs * bin_centers, dim=-1)

                gripper_probs = F.softmax(gripper_logits, dim=-1)
                p1 = gripper_probs[..., 1]
                pred_gripper = -1.0 + 2.0 * p1
                
                pred_action_continuous = torch.cat([pred_pose, pred_gripper], dim=-1)

        if self.dct_loss_weight > 0 and pred_action_continuous is not None:
             loss_dct = self._compute_dct_loss(pred_action_continuous.float(), actions.float())
             loss = loss + self.dct_loss_weight * loss_dct
        
        if self.future_image_loss_weight > 0:
            loss = loss + self.future_image_loss_weight * loss_img
        return loss

    def _forward_regression(self, input_ids, attention_mask, actions, proprioception, history_actions, proprio_attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw, future_images=None, token_type_ids=None):
        connector_out, hidden_states = self.get_vlm_condition(
            input_ids, attention_mask, proprioception=proprioception, proprio_attention_mask=proprio_attention_mask,
            pixel_values=pixel_values, pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw,
            token_type_ids=token_type_ids
        )

        loss_img = 0.0
        gen_hidden_states = None
        if self.enable_future_image_loss and future_images is not None:
             loss_img, gen_hidden_states = self._compute_gen_loss_and_feats(future_images, hidden_states)
        
        policy_history = history_actions if self.use_action_input_policy else None
        
        if self.condition_type in ["tight", "soft"]:
             if self.enable_future_image_loss:
                 pred_actions = self.action_head(hidden_states, history_actions=policy_history, gen_hidden_states=gen_hidden_states)
             else:
                 pred_actions = self.action_head(hidden_states, history_actions=policy_history)
        elif self.condition_type == "loose":
             cond_input = connector_out.mean(dim=1)
             pred_actions = self.action_head(cond_input, history_actions=policy_history)
        else:
             raise ValueError(f"Unknown condition type: {self.condition_type}")

        if actions.ndim == 2: actions = actions.unsqueeze(1)
        loss = F.mse_loss(pred_actions, actions)

        if self.dct_loss_weight > 0:
            loss_dct = self._compute_dct_loss(pred_actions.float(), actions.float())
            loss = loss + self.dct_loss_weight * loss_dct
            
        if self.future_image_loss_weight > 0:
            loss = loss + self.future_image_loss_weight * loss_img
        return loss

    def _forward_diffusion(self, input_ids, attention_mask, actions, proprioception, history_actions, proprio_attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw, future_images=None, token_type_ids=None):
        connector_out, hidden_states = self.get_vlm_condition(
            input_ids, attention_mask, proprioception=proprioception, proprio_attention_mask=proprio_attention_mask,
            pixel_values=pixel_values, pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw,
            token_type_ids=token_type_ids
        )

        loss_img = 0.0
        gen_hidden_states = None
        if self.enable_future_image_loss and future_images is not None:
             loss_img, gen_hidden_states = self._compute_gen_loss_and_feats(future_images, hidden_states)
        
        if actions.ndim == 2: actions = actions.unsqueeze(1)
        noise = torch.randn_like(actions)
        B = actions.shape[0]
        
        if self.scheduler_type == "flow_match":
            sigmas = torch.rand((B,), device=actions.device)
            sigmas_expanded = sigmas.view(B, *([1] * (actions.ndim - 1)))
            noisy_actions = (1.0 - sigmas_expanded) * actions + sigmas_expanded * noise
            noisy_actions = noisy_actions.to(dtype=actions.dtype)
            timesteps = sigmas * self.noise_scheduler.config.num_train_timesteps
            target = noise - actions
        else:
            timesteps = torch.randint(0, self.noise_scheduler.config.num_train_timesteps, (B,), device=actions.device).long()
            noisy_actions = self.noise_scheduler.add_noise(actions, noise, timesteps)
            target = noise
            
        policy_history = history_actions if self.use_action_input_policy else None

        if self.condition_type in ["tight", "soft"]:
            if self.enable_future_image_loss:
                pred = self.action_head(noisy_actions, timesteps, hidden_states, history_actions=policy_history, gen_hidden_states=gen_hidden_states)
            else:
                pred = self.action_head(noisy_actions, timesteps, hidden_states, history_actions=policy_history)
        elif self.condition_type == "loose":
            cond_input = connector_out.mean(dim=1)
            pred = self.action_head(noisy_actions, timesteps, cond_input, history_actions=policy_history)
        else:
             raise ValueError(f"Unknown condition type: {self.condition_type}")
        
        loss = F.mse_loss(pred, target)

        if self.dct_loss_weight > 0:
            pred_x_start = None
            if self.scheduler_type == "flow_match":
                 pred_x_start = noisy_actions - sigmas_expanded * pred
            elif self.scheduler_type == "ddim":
                 def view_right(t):
                    while t.ndim < pred.ndim:
                        t = t.unsqueeze(-1)
                    return t
                 alphas_cumprod = self.noise_scheduler.alphas_cumprod.to(device=pred.device, dtype=pred.dtype)
                 alpha_prod_t = alphas_cumprod[timesteps]
                 pred_x_start = (noisy_actions - view_right((1 - alpha_prod_t).sqrt()) * pred) / view_right(alpha_prod_t.sqrt())

            if pred_x_start is not None:
                 loss_dct = self._compute_dct_loss(pred_x_start.float(), actions.float())
                 loss = loss + self.dct_loss_weight * loss_dct
            
        if self.future_image_loss_weight > 0:
            loss = loss + self.future_image_loss_weight * loss_img
        return loss

    @torch.no_grad()
    def predict_action(self, input_ids, attention_mask, proprioception=None, history_actions=None, proprio_attention_mask=None, pixel_values=None, pixel_values_videos=None, image_grid_thw=None, video_grid_thw=None, token_type_ids=None):
        B = input_ids.shape[0]

        connector_out, hidden_states = self.get_vlm_condition(
            input_ids, attention_mask,
            proprioception=proprioception,
            proprio_attention_mask=proprio_attention_mask,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            token_type_ids=token_type_ids
        )
        
        policy_history = history_actions if self.use_action_input_policy else None
        gen_hidden_states = None
        if self.enable_future_image_loss and self.condition_type in ["tight", "soft"]:
             num_img_tokens = 256
             gen_context = hidden_states
             gen_hidden_states = None
             if self.use_incremental_gen:
                 # O(n) incremental decode -> token IDs + per-layer hidden states
                 # in one pass (equivalent to the AR loop + the line-887 forward).
                 res = self.generator.generate_incremental(gen_context, num_img_tokens)
                 if res is not None:
                     _, gen_hidden_states = res
             if gen_hidden_states is None:
                 # full-recompute fallback (original O(n^2) path)
                 curr_ids = torch.zeros((B, 1), dtype=torch.long, device=input_ids.device)
                 for _ in range(num_img_tokens):
                     logits, _ = self.generator(curr_ids, gen_context)
                     next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
                     curr_ids = torch.cat([curr_ids, next_token], dim=1)
                 gen_input = curr_ids[:, :-1]
                 _, gen_hidden_states = self.generator(gen_input, gen_context)

        if self.loss_type == "regression":
            if self.condition_type in ["tight", "soft"]:
                if self.enable_future_image_loss:
                    action = self.action_head(hidden_states, history_actions=policy_history, gen_hidden_states=gen_hidden_states)
                else:
                    action = self.action_head(hidden_states, history_actions=policy_history)
            elif self.condition_type == "loose":
                cond_input = connector_out.mean(dim=1)
                action = self.action_head(cond_input, history_actions=policy_history)
            if action.ndim == 2 and self.num_actions > 1:
                action = action.view(action.shape[0], self.num_actions, self.action_dim)
            
            return action.to(dtype=self.lmm.dtype)

        elif self.loss_type == "classification":
            if self.condition_type in ["tight", "soft"]:
                if self.enable_future_image_loss:
                    logits = self.action_head(hidden_states, history_actions=policy_history, gen_hidden_states=gen_hidden_states)
                else:
                    logits = self.action_head(hidden_states, history_actions=policy_history)
            else:
                cond_input = connector_out.mean(dim=1)
                logits = self.action_head(cond_input, history_actions=policy_history)

            if self.action_vqvae is not None:
                indices = torch.argmax(logits, dim=-1)  # (B, T, Latent_Codes)
                action = self.action_vqvae.decode_indices(indices)
                return action.to(dtype=self.lmm.dtype)

            else:
                pose_logits = logits[:, :, :self.action_dim - 1, :]
                gripper_logits = logits[:, :, -1:, :2]
                pose_idx = torch.argmax(pose_logits, dim=-1)
                gripper_idx = torch.argmax(gripper_logits, dim=-1)
                pose_pred = (pose_idx.float() / (self.num_bins - 1)) * 2 - 1
                gripper_pred = gripper_idx.float() * 2 - 1
                action = torch.cat([pose_pred, gripper_pred], dim=-1).to(dtype=self.lmm.dtype)
                return action

        elif self.loss_type == "diffusion":
            action = torch.randn(B, self.num_actions, self.action_dim, device=input_ids.device).to(self.lmm.dtype)
            self.noise_scheduler.set_timesteps(self.num_inference_timesteps)
            
            for t in self.noise_scheduler.timesteps:
                timesteps = torch.full((B,), t, device=input_ids.device)
                if self.scheduler_type != "flow_match": timesteps = timesteps.long()
                if self.condition_type in ["tight", "soft"]:
                    if self.enable_future_image_loss:
                        output = self.action_head(action, timesteps, hidden_states, history_actions=policy_history, gen_hidden_states=gen_hidden_states)
                    else:
                        output = self.action_head(action, timesteps, hidden_states, history_actions=policy_history)
                else:
                    cond_input = connector_out.mean(dim=1)
                    output = self.action_head(action, timesteps, cond_input, history_actions=policy_history)
                
                action = self.noise_scheduler.step(output, t, action).prev_sample
                action = action.to(dtype=self.lmm.dtype)
            
            return action
        
        else:
            raise ValueError(f"Unknown loss type: {self.loss_type}")

    @torch.no_grad()
    def predict_image(self, input_ids, attention_mask, proprioception=None, history_actions=None, proprio_attention_mask=None, pixel_values=None, pixel_values_videos=None, image_grid_thw=None, video_grid_thw=None, max_new_tokens=1024, token_type_ids=None):
        _, hidden_states = self.get_vlm_condition(
            input_ids, attention_mask,
            proprioception=proprioception,
            proprio_attention_mask=proprio_attention_mask,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            token_type_ids=token_type_ids
        )
        gen_vlm_ctx = hidden_states
        
        curr_ids = torch.zeros((input_ids.shape[0], 1), dtype=torch.long, device=input_ids.device)
        
        for _ in range(max_new_tokens):
            logits, _ = self.generator(curr_ids, gen_vlm_ctx)
            next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
            curr_ids = torch.cat([curr_ids, next_token], dim=1)
            
        generated_tokens = curr_ids[:, 1:]
        H_latent = int(generated_tokens.shape[1]**0.5)
        decoded_images = self.vq_model.decode_code(generated_tokens, shape=(input_ids.shape[0], H_latent, H_latent))
        return decoded_images

if __name__ == "__main__":
    print("Testing VLANeXt Model...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16
    
    # Initialize Model (Minimal Config)
    model = VLANeXt(
        lmm_path="Qwen/Qwen3-VL-2B-Instruct",
        action_dim=7, num_actions=4, num_history=2,
        backbone_mode="finetune", gradient_checkpointing=False
    ).to(device, dtype)
    processor = model.processor

    def run_test(modality="image"):
        print(f"\n=== Testing {modality.capitalize()} ===")
        B = 2
        # Dummy Data
        img = Image.new('RGB', (64, 64), color='red')
        media = [img] * B if modality == "image" else [[img]*8] * B
        content_key = "image" if modality == "image" else "video"
        
        # Process
        msgs = [[{"role": "user", "content": [{"type": content_key, content_key: m}, {"type": "text", "text": "Task."}]}] for m in media]
        texts = [processor.apply_chat_template(m, tokenize=False, add_generation_prompt=True) for m in msgs]
        inputs = processor(text=texts, **{f"{modality}s": media}, padding=True, return_tensors="pt")
        
        # Move to device & cast
        inputs = {k: v.to(device) for k, v in inputs.items()}
        for k in ["pixel_values", "pixel_values_videos"]:
            if k in inputs: inputs[k] = inputs[k].to(dtype)
            
        # Filter valid args for forward
        valid_keys = {"input_ids", "attention_mask", "pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw"}
        fwd_args = {k: v for k, v in inputs.items() if k in valid_keys}

        # Tensors
        act_gt = torch.randn(B, 4, 7, device=device, dtype=dtype)
        proprio = torch.randn(B, 2, 7, device=device, dtype=dtype)
        hist_act = torch.randn(B, 2, 7, device=device, dtype=dtype)

        # Tests
        print(f"Action Gen Loss: {model(actions=act_gt, proprioception=proprio, history_actions=hist_act, **fwd_args).item():.4f}")
        print(f"Action Pred Shape: {model.predict_action(proprioception=proprio, history_actions=hist_act, **fwd_args).shape}")

    run_test("image")
    run_test("video")
    print("\nTest Passed!")