vllm-omni/vllm_omni/diffusion/models/internvla_a1/model_internvla_a1.py

from __future__ import annotations

import math
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal

import torch
import torch._dynamo as dynamo
import torch.nn.functional as F
from einops import rearrange
from safetensors import safe_open
from torch import Tensor, nn
from transformers import Qwen3VLProcessor
from transformers.models.auto import CONFIG_MAPPING
from vllm.model_executor.models.utils import AutoWeightsLoader

from .adapter_qwen3_vl import (
    Qwen3VLForConditionalGeneration,
    Qwen3VLTextModel,
    apply_rotary_pos_emb,
    eager_attention_forward,
)
from .config import (
    DEFAULT_COSMOS_DIR,
    DEFAULT_COSMOS_REPO,
    DEFAULT_QWEN3_VL_MODEL,
    OBS_IMAGES,
    OBS_PREFIX,
    OBS_STATE,
    OBS_TASK,
    OPENPI_ATTENTION_MASK_VALUE,
    InternVLAA1Config,
)
from .model_cosmos import ImageTokenizer


def get_safe_dtype(target_dtype: torch.dtype, device_type: str) -> torch.dtype:
    if device_type == "mps" and target_dtype == torch.float64:
        return torch.float32
    if device_type == "cpu" and target_dtype == torch.bfloat16:
        return torch.float32
    return target_dtype


def create_sinusoidal_pos_embedding(
    time: torch.Tensor,
    dimension: int,
    min_period: float,
    max_period: float,
    device: torch.device,
) -> Tensor:
    if dimension % 2 != 0:
        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
    dtype = get_safe_dtype(torch.float64, device.type)
    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
    period = min_period * (max_period / min_period) ** fraction
    scaling_factor = 1.0 / period * 2 * math.pi
    sin_input = scaling_factor[None, :] * time[:, None]
    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)


def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
    cumsum = torch.cumsum(att_masks, dim=1)
    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
    return att_2d_masks & pad_2d_masks


def pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
    if vector.shape[-1] >= new_dim:
        return vector
    return F.pad(vector, (0, new_dim - vector.shape[-1]))


def resize_with_pad(images: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
    target_h, target_w = size
    if images.ndim != 4:
        raise ValueError(f"Expected [T, C, H, W], got {tuple(images.shape)}")
    _, _, src_h, src_w = images.shape
    scale = min(target_h / src_h, target_w / src_w)
    resized_h = max(1, int(round(src_h * scale)))
    resized_w = max(1, int(round(src_w * scale)))
    resized = F.interpolate(images, size=(resized_h, resized_w), mode="bilinear", align_corners=False)
    pad_h = target_h - resized_h
    pad_w = target_w - resized_w
    pad_top = pad_h // 2
    pad_bottom = pad_h - pad_top
    pad_left = pad_w // 2
    pad_right = pad_w - pad_left
    return F.pad(resized, (pad_left, pad_right, pad_top, pad_bottom))


class Qwen3VLInputBuilder:
    def __init__(
        self,
        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
        max_length: int = 48,
        spatial_merge_size: int = 2,
    ) -> None:
        processor_model_name = os.getenv("INTERNVLA_A1_PROCESSOR_DIR", processor_model_name)
        local_path = Path(processor_model_name).expanduser()
        if local_path.exists():
            self.processor = Qwen3VLProcessor.from_pretrained(str(local_path), local_files_only=True)
        else:
            if os.getenv("INTERNVLA_A1_PROCESSOR_DIR"):
                raise FileNotFoundError(f"INTERNVLA_A1_PROCESSOR_DIR points to a missing path: {local_path}")
            self.processor = Qwen3VLProcessor.from_pretrained(processor_model_name)
        self.max_length = max_length
        self.spatial_merge_size = spatial_merge_size

    def build(self, camera_images: list[torch.Tensor], task: str) -> dict[str, torch.Tensor]:
        input_ids: list[int] = []
        attention_mask: list[int] = []
        pixel_values: list[torch.Tensor] = []
        image_grid_thw: list[torch.Tensor] = []

        for image_history in camera_images:
            current_image = image_history[1]
            image_inputs = self.processor.image_processor(current_image, do_rescale=False)
            pixel_values.append(image_inputs.pixel_values)
            image_grid_thw.append(image_inputs.image_grid_thw)
            num_img_token = int(torch.prod(image_inputs.image_grid_thw[0]).item()) // (self.spatial_merge_size**2)
            input_ids.extend(
                [self.processor.vision_start_token_id]
                + [self.processor.image_token_id] * num_img_token
                + [self.processor.vision_end_token_id]
            )
            attention_mask.extend([1] * (num_img_token + 2))

        text_inputs = self.processor.tokenizer(
            task,
            max_length=self.max_length,
            padding="max_length",
            padding_side="right",
            truncation=True,
        )
        input_ids.extend(text_inputs.input_ids)
        attention_mask.extend(text_inputs.attention_mask)

        return {
            f"{OBS_PREFIX}pixel_values": torch.cat(pixel_values, dim=0),
            f"{OBS_PREFIX}image_grid_thw": torch.cat(image_grid_thw, dim=0),
            f"{OBS_PREFIX}input_ids": torch.tensor(input_ids, dtype=torch.long),
            f"{OBS_PREFIX}attention_mask": torch.tensor(attention_mask, dtype=torch.long),
        }


@dataclass
class SuffixStaticContext:
    state_emb: torch.Tensor
    full_att_2d_masks_4d: torch.Tensor
    position_ids: torch.Tensor


def compute_layer_complete(
    layer_idx: int,
    inputs_embeds: list[torch.Tensor],
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    und_expert: nn.Module,
    gen_expert: nn.Module,
    act_expert: nn.Module,
) -> list[torch.Tensor]:
    models = [und_expert.language_model, gen_expert, act_expert]
    query_states = []
    key_states = []
    value_states = []
    for i, hidden_states in enumerate(inputs_embeds):
        layer = models[i].layers[layer_idx]
        hidden_states = layer.input_layernorm(hidden_states)
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
        if layer.self_attn.q_proj.weight.dtype == torch.bfloat16:
            hidden_states = hidden_states.to(dtype=torch.bfloat16)
        query_state = layer.self_attn.q_norm(layer.self_attn.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        key_state = layer.self_attn.k_norm(layer.self_attn.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        query_states.append(query_state)
        key_states.append(key_state)
        value_states.append(value_state)

    query_states = torch.cat(query_states, dim=2)
    key_states = torch.cat(key_states, dim=2)
    value_states = torch.cat(value_states, dim=2)
    dummy_tensor = torch.zeros(
        query_states.shape[0],
        query_states.shape[2],
        query_states.shape[-1],
        device=query_states.device,
        dtype=query_states.dtype,
    )
    cos, sin = und_expert.model.language_model.rotary_emb(dummy_tensor, position_ids)
    query_states, key_states = apply_rotary_pos_emb(
        query_states,
        key_states,
        cos,
        sin,
        unsqueeze_dim=1,
    )
    scaling = und_expert.language_model.layers[layer_idx].self_attn.scaling
    att_output, _ = eager_attention_forward(
        und_expert.language_model.layers[layer_idx].self_attn,
        query_states,
        key_states,
        value_states,
        attention_mask,
        scaling,
    )

    head_dim = und_expert.language_model.layers[layer_idx].self_attn.head_dim
    num_attention_heads = und_expert.language_model.layers[layer_idx].self_attn.config.num_attention_heads
    batch_size = query_states.shape[0]
    att_output = att_output.reshape(batch_size, -1, num_attention_heads * head_dim)

    outputs_embeds = []
    start_pos = 0
    for i, hidden_states in enumerate(inputs_embeds):
        layer = models[i].layers[layer_idx]
        end_pos = start_pos + hidden_states.shape[1]
        att_chunk = att_output[:, start_pos:end_pos]
        if att_chunk.dtype != layer.self_attn.o_proj.weight.dtype:
            att_chunk = att_chunk.to(layer.self_attn.o_proj.weight.dtype)
        out_emb = layer.self_attn.o_proj(att_chunk)
        out_emb = out_emb + hidden_states
        residual = out_emb
        out_emb = layer.post_attention_layernorm(out_emb)
        if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
            out_emb = out_emb.to(dtype=torch.bfloat16)
        out_emb = layer.mlp(out_emb)
        out_emb = out_emb + residual
        outputs_embeds.append(out_emb)
        start_pos = end_pos
    return outputs_embeds


class QwenConfig:
    def __init__(
        self,
        head_dim: int,
        hidden_size: int,
        intermediate_size: int,
        num_attention_heads: int,
        num_hidden_layers: int,
        num_key_value_heads: int,
    ) -> None:
        self.head_dim = head_dim
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.num_key_value_heads = num_key_value_heads


def get_qwen_config(variant: str) -> QwenConfig:
    num_hidden_layers = int(variant.split("_")[-1][:-1])
    if variant.startswith("qwen3_vl"):
        return QwenConfig(128, 2048, 6144, 16, num_hidden_layers, 8)
    if variant.startswith("qwen3"):
        return QwenConfig(128, 1024, 3072, 16, num_hidden_layers, 8)
    raise ValueError(f"Unknown variant: {variant}")


def resolve_cosmos_checkpoint_paths(
    *,
    encoder_path: str | Path | None = None,
    decoder_path: str | Path | None = None,
) -> tuple[Path, Path]:
    encoder_override = encoder_path or os.getenv("INTERNVLA_A1_COSMOS_ENCODER_PATH")
    decoder_override = decoder_path or os.getenv("INTERNVLA_A1_COSMOS_DECODER_PATH")
    default_encoder_name = "encoder.safetensors"
    default_decoder_name = "decoder.safetensors"
    checkpoint_enc = (
        Path(encoder_override).expanduser()
        if encoder_override is not None
        else DEFAULT_COSMOS_DIR / default_encoder_name
    )
    checkpoint_dec = (
        Path(decoder_override).expanduser()
        if decoder_override is not None
        else DEFAULT_COSMOS_DIR / default_decoder_name
    )

    missing = [str(path) for path in (checkpoint_enc, checkpoint_dec) if not path.exists()]
    if missing:
        missing_display = ", ".join(missing)
        raise FileNotFoundError(
            "InternVLA-A1 requires local Cosmos tokenizer checkpoints. "
            f"Missing: {missing_display}. "
            f"Download {DEFAULT_COSMOS_REPO} and set INTERNVLA_A1_COSMOS_DIR to that directory, "
            "or point INTERNVLA_A1_COSMOS_ENCODER_PATH / INTERNVLA_A1_COSMOS_DECODER_PATH "
            "to explicit checkpoint files."
        )

    return checkpoint_enc, checkpoint_dec


class Qwen3VLWithExpertModel(nn.Module):
    def __init__(
        self,
        vlm_config: QwenConfig,
        action_expert_config: QwenConfig,
        precision: Literal["bfloat16", "float32"] = "bfloat16",
    ) -> None:
        super().__init__()

        vlm_config_hf = CONFIG_MAPPING["qwen3_vl"]()
        vlm_config_hf.text_config.hidden_size = vlm_config.hidden_size
        vlm_config_hf.text_config.intermediate_size = vlm_config.intermediate_size
        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_attention_heads
        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
        vlm_config_hf.text_config.num_hidden_layers = vlm_config.num_hidden_layers
        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_key_value_heads
        vlm_config_hf.text_config.max_position_embeddings = 262144
        vlm_config_hf.text_config.rope_scaling = {
            "mrope_interleaved": True,
            "mrope_section": [24, 20, 20],
            "rope_type": "default",
        }
        vlm_config_hf.text_config.tie_word_embeddings = True
        vlm_config_hf.tie_word_embeddings = True
        vlm_config_hf.vision_config.deepstack_visual_indexes = [5, 11, 17]
        vlm_config_hf.vision_config.depth = 24
        vlm_config_hf.vision_config.hidden_size = 1024
        vlm_config_hf.vision_config.intermediate_size = 4096
        vlm_config_hf.vision_config.out_hidden_size = 2048

        self.und_expert = Qwen3VLForConditionalGeneration(config=vlm_config_hf)

        gen_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]()
        gen_expert_config_hf.head_dim = action_expert_config.head_dim
        gen_expert_config_hf.hidden_size = action_expert_config.hidden_size
        gen_expert_config_hf.intermediate_size = action_expert_config.intermediate_size
        gen_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads
        gen_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers
        gen_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads
        gen_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings
        gen_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling
        self.gen_expert = Qwen3VLTextModel(config=gen_expert_config_hf)
        self.gen_expert.embed_tokens = None
        self.gen_expert.lm_head = None

        act_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]()
        act_expert_config_hf.head_dim = action_expert_config.head_dim
        act_expert_config_hf.hidden_size = action_expert_config.hidden_size
        act_expert_config_hf.intermediate_size = action_expert_config.intermediate_size
        act_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads
        act_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers
        act_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads
        act_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings
        act_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling
        self.act_expert = Qwen3VLTextModel(config=act_expert_config_hf)
        self.act_expert.embed_tokens = None
        self.act_expert.lm_head = None

        self.to_selected_precision(precision)

    def to_selected_precision(self, precision: Literal["bfloat16", "float32"]) -> None:
        if precision == "bfloat16":
            self.to(dtype=torch.bfloat16)
        elif precision == "float32":
            # The fp32 path keeps the whole module in full precision, so no
            # selective post-conversion fixups are required.
            self.to(dtype=torch.float32)
            return
        else:
            raise ValueError(f"Unsupported precision: {precision}")

        # Match the original runtime numerics by keeping normalization layers
        # in fp32 even when the rest of the module runs in bf16.
        keep_fp32 = ["input_layernorm", "post_attention_layernorm", "model.norm"]
        for name, param in self.named_parameters():
            if any(key in name for key in keep_fp32):
                param.data = param.data.to(dtype=torch.float32)

    def forward(
        self,
        attention_mask: torch.Tensor | None,
        position_ids: torch.LongTensor | None,
        past_key_values: Any,
        inputs_embeds: list[torch.Tensor | None],
        use_cache: bool,
    ) -> tuple[list[torch.Tensor | None], Any]:
        if inputs_embeds[1] is None and inputs_embeds[2] is None:
            prefix_output = self.und_expert.language_model.forward(
                inputs_embeds=inputs_embeds[0],
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
            )
            return [prefix_output.last_hidden_state, None, None], prefix_output.past_key_values

        if inputs_embeds[0] is None and inputs_embeds[2] is None:
            middle_output = self.gen_expert.forward(
                inputs_embeds=inputs_embeds[1],
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
            )
            return [None, middle_output.last_hidden_state, None], middle_output.past_key_values

        if inputs_embeds[0] is None and inputs_embeds[1] is None:
            suffix_output = self.act_expert.forward(
                inputs_embeds=inputs_embeds[2],
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
            )
            return [None, None, suffix_output.last_hidden_state], None

        models = [self.und_expert.language_model, self.gen_expert, self.act_expert]
        stacked_inputs = [inputs_embeds[0], inputs_embeds[1], inputs_embeds[2]]
        for layer_idx in range(self.und_expert.config.text_config.num_hidden_layers):
            stacked_inputs = compute_layer_complete(
                layer_idx,
                stacked_inputs,
                attention_mask,
                position_ids,
                und_expert=self.und_expert,
                gen_expert=self.gen_expert,
                act_expert=self.act_expert,
            )
        outputs = [models[i].norm(stacked_inputs[i]) for i in range(3)]
        return outputs, None


class InternVLAA1(nn.Module):
    def __init__(
        self,
        config: InternVLAA1Config,
        *,
        cosmos_encoder_path: str | Path | None = None,
        cosmos_decoder_path: str | Path | None = None,
    ) -> None:
        super().__init__()
        self.config = config

        vlm_config = get_qwen_config(config.qwen3_vl_variant)
        action_expert_config = get_qwen_config(config.action_expert_variant)
        self.qwen3_vl_with_expert = Qwen3VLWithExpertModel(vlm_config, action_expert_config, precision=config.dtype)

        cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths(
            encoder_path=cosmos_encoder_path,
            decoder_path=cosmos_decoder_path,
        )
        self.cosmos = ImageTokenizer(
            checkpoint_enc=str(cosmos_encoder_path),
            checkpoint_dec=str(cosmos_decoder_path),
            device=config.device,
        )

        hidden_size = action_expert_config.hidden_size
        vae_dim = 16
        ds = config.scale_factor
        self.cosmos_in_proj = nn.Conv2d(vae_dim, hidden_size, kernel_size=1, stride=1, padding=0)
        self.downsample_conv = nn.Conv2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0)
        self.upsample_conv = nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0)
        self.cosmos_out_proj = nn.Linear(hidden_size, vae_dim)
        self.cosmos_out_layer_norm = nn.LayerNorm(hidden_size)

        self.action_in_proj = nn.Linear(config.max_action_dim, hidden_size)
        self.action_out_proj = nn.Linear(hidden_size, config.max_action_dim)
        self.state_proj = nn.Linear(config.max_state_dim, hidden_size)
        self.action_time_mlp_in = nn.Linear(2 * hidden_size, hidden_size)
        self.action_time_mlp_out = nn.Linear(hidden_size, hidden_size)

        if config.compile_model:
            torch.set_float32_matmul_precision("high")
            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)

        self.cosmos.eval()
        for param in self.cosmos.parameters():
            param.requires_grad = False

        self.set_attention_implementation(config.attn_implementation)

    def set_attention_implementation(self, attn_implementation: str) -> None:
        self.config.attn_implementation = attn_implementation
        self.qwen3_vl_with_expert.und_expert.config.text_config._attn_implementation = attn_implementation
        self.qwen3_vl_with_expert.und_expert.language_model.config._attn_implementation = attn_implementation
        self.qwen3_vl_with_expert.gen_expert.config._attn_implementation = attn_implementation
        self.qwen3_vl_with_expert.act_expert.config._attn_implementation = attn_implementation

    def _prepare_attention_masks_4d(
        self,
        att_2d_masks: torch.Tensor,
        *,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        att_2d_masks_4d = att_2d_masks[:, None, :, :]
        bias_dtype = torch.float32 if self.config.attn_implementation == "eager" else dtype
        valid_bias = torch.zeros((), dtype=bias_dtype, device=att_2d_masks.device)
        invalid_bias = torch.tensor(
            OPENPI_ATTENTION_MASK_VALUE,
            dtype=bias_dtype,
            device=att_2d_masks.device,
        )
        return torch.where(att_2d_masks_4d, valid_bias, invalid_bias)

    def sample_noise(self, shape: tuple[int, ...], device: torch.device) -> torch.Tensor:
        return torch.normal(0.0, 1.0, shape, dtype=torch.float32, device=device)

    @dynamo.disable
    def embed_prefix(
        self,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.Tensor,
        lang_tokens: torch.Tensor,
        lang_masks: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        image_token_id = self.qwen3_vl_with_expert.und_expert.config.image_token_id
        pixel_values = pixel_values.view(-1, pixel_values.shape[-1])
        image_grid_thw = image_grid_thw.view(-1, 3)
        image_embs, _ = self.qwen3_vl_with_expert.und_expert.visual(pixel_values, image_grid_thw)

        embs = self.qwen3_vl_with_expert.und_expert.get_input_embeddings()(lang_tokens)
        batch_size, seq_len, hidden_dim = embs.shape
        embs = embs.view(-1, hidden_dim)
        lang_tokens_flat = lang_tokens.view(-1)
        embs[lang_tokens_flat == image_token_id] = image_embs
        embs = embs.view(batch_size, seq_len, hidden_dim)

        pad_masks = lang_masks.to(torch.bool)
        att_masks = torch.zeros_like(pad_masks, dtype=torch.bool, device=pad_masks.device)
        return embs, pad_masks, att_masks

    def get_cosmos_features(self, images: torch.Tensor) -> torch.Tensor:
        shape = images.shape[:-3]
        channels, height, width = images.shape[-3:]
        images = images.reshape(-1, channels, height, width)
        if (height, width) != (256, 256):
            images = F.interpolate(images, size=(256, 256), mode="bilinear", align_corners=False)
        images = images * 2 - 1
        features = self.cosmos.encode(images)
        channels, height, width = features.shape[-3:]
        return features.view(*shape, channels, height, width)

    def embed_middle(
        self,
        images: torch.Tensor,
        img_masks: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        device = images.device
        batch_size, n_view, timesteps = images.shape[:3]
        features = self.get_cosmos_features(images)
        features = rearrange(features, "b n t c h w -> (b n t) c h w")
        features = self.cosmos_in_proj(features)
        features = self.downsample_conv(features)
        features = rearrange(features, "(b n t) c h w -> b n t c h w", b=batch_size, n=n_view, t=timesteps)

        _, _, _, _, height, width = features.shape
        embs = rearrange(features, "b n t c h w -> b (n t h w) c")
        pad_masks = torch.zeros((batch_size, n_view, timesteps, height, width), dtype=torch.bool, device=device)
        pad_masks[img_masks] = True
        pad_masks = rearrange(pad_masks, "b n t h w -> b (n t h w)")
        att_masks = torch.tensor([1] + [0] * (embs.shape[1] - 1), dtype=torch.bool, device=device)
        att_masks = att_masks[None, :].expand(batch_size, att_masks.shape[0])
        return embs, pad_masks, att_masks

    def prepare_suffix_static_context(
        self,
        state: torch.Tensor,
        prefix_pad_masks: torch.Tensor,
        max_prefix_position_ids: torch.Tensor,
    ) -> SuffixStaticContext:
        if self.state_proj.weight.dtype == torch.float32:
            state = state.to(torch.float32)

        state_emb = self.state_proj(state)
        device = state_emb.device
        batch_size = state_emb.shape[0]
        suffix_len = self.config.chunk_size + 1
        prefix_len = prefix_pad_masks.shape[1]

        suffix_pad_masks = torch.ones(batch_size, suffix_len, dtype=torch.bool, device=device)
        suffix_att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=torch.bool, device=device)
        suffix_att_masks = suffix_att_masks[None, :].expand(batch_size, suffix_len)
        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
            torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2),
            dtype=state_emb.dtype,
        )
        position_ids = (
            torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device)
            .repeat(3, 1, 1)
            .to(max_prefix_position_ids)
            + max_prefix_position_ids
        )
        return SuffixStaticContext(
            state_emb=state_emb,
            full_att_2d_masks_4d=full_att_2d_masks_4d,
            position_ids=position_ids,
        )

    def embed_suffix(
        self,
        state: torch.Tensor,
        noisy_actions: torch.Tensor,
        timestep: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if self.state_proj.weight.dtype == torch.float32:
            state = state.to(torch.float32)

        state_emb = self.state_proj(state)
        batch_size = state_emb.shape[0]
        device = state_emb.device
        state_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=device)

        time_emb = create_sinusoidal_pos_embedding(
            timestep,
            self.action_in_proj.out_features,
            min_period=self.config.min_period,
            max_period=self.config.max_period,
            device=timestep.device,
        ).to(dtype=timestep.dtype)
        action_emb = self.action_in_proj(noisy_actions)
        time_emb = time_emb[:, None, :].expand_as(action_emb)
        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
        action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb)))
        action_mask = torch.ones(batch_size, action_time_emb.shape[1], dtype=torch.bool, device=device)

        embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1)
        pad_masks = torch.cat([state_mask, action_mask], dim=1)
        att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=embs.dtype, device=device)
        att_masks = att_masks[None, :].expand(batch_size, self.config.chunk_size + 1)
        return embs, pad_masks, att_masks

    def get_position_ids(
        self,
        lang_tokens: torch.Tensor,
        image_grid_thw: torch.Tensor | None,
        pad_masks: torch.Tensor,
    ) -> tuple[torch.Tensor, Any]:
        seq_len = lang_tokens.shape[1]
        padded_lang_tokens = torch.ones_like(pad_masks).to(lang_tokens) * 777
        padded_lang_tokens[:, :seq_len] = lang_tokens
        attention_mask = pad_masks.to(lang_tokens)
        if image_grid_thw is not None:
            image_grid_thw = image_grid_thw.view(-1, 3)
        return self.qwen3_vl_with_expert.und_expert.model.get_rope_index(
            padded_lang_tokens,
            image_grid_thw,
            attention_mask=attention_mask,
        )

    @torch.no_grad()
    def sample_actions(
        self,
        images: torch.Tensor,
        img_masks: torch.Tensor,
        pixel_values: torch.Tensor,
        image_grid_thw: torch.Tensor,
        lang_tokens: torch.Tensor,
        lang_masks: torch.Tensor,
        state: torch.Tensor,
        *,
        noise: torch.Tensor | None = None,
        num_steps: int | None = None,
        decode_image: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if num_steps is None:
            num_steps = self.config.num_inference_steps
        batch_size = state.shape[0]
        device = state.device
        dtype = state.dtype
        if noise is None:
            noise = self.sample_noise((batch_size, self.config.chunk_size, self.config.max_action_dim), device)

        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
            pixel_values,
            image_grid_thw,
            lang_tokens,
            lang_masks,
        )
        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
        prefix_position_ids, _ = self.get_position_ids(lang_tokens, image_grid_thw, prefix_pad_masks)
        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(
            prefix_att_2d_masks,
            dtype=prefix_embs.dtype,
        )
        _, past_key_values = self.qwen3_vl_with_expert.forward(
            attention_mask=prefix_att_2d_masks_4d,
            position_ids=prefix_position_ids,
            past_key_values=None,
            inputs_embeds=[prefix_embs, None, None],
            use_cache=True,
        )
        max_prefix_position_ids = prefix_position_ids.max(dim=-1, keepdim=True).values

        middle_embs, middle_pad_masks, middle_att_masks = self.embed_middle(images[:, :, :2], img_masks)
        middle_len = middle_pad_masks.shape[1]
        prefix_len = prefix_pad_masks.shape[1]
        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, middle_len, prefix_len)
        middle_att_2d_masks = make_att_2d_masks(middle_pad_masks, middle_att_masks)
        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, middle_att_2d_masks], dim=2)
        middle_position_ids = (
            torch.arange(1, middle_len + 1, device=max_prefix_position_ids.device)
            .repeat(3, 1, 1)
            .to(max_prefix_position_ids)
            + max_prefix_position_ids
        )
        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
            full_att_2d_masks,
            dtype=middle_embs.dtype,
        )
        (_, middle_out, _), past_key_values = self.qwen3_vl_with_expert.forward(
            attention_mask=full_att_2d_masks_4d,
            position_ids=middle_position_ids,
            past_key_values=past_key_values,
            inputs_embeds=[None, middle_embs, None],
            use_cache=True,
        )

        max_position_ids = middle_position_ids.max(dim=-1, keepdim=True).values
        curr_pad_masks = torch.cat([prefix_pad_masks, middle_pad_masks], dim=1)
        dt = torch.tensor(-1.0 / num_steps, dtype=torch.float32, device=device)
        x_t = noise
        time = torch.tensor(1.0, dtype=torch.float32, device=device)
        suffix_static = None
        if self.config.enable_suffix_static_context_optimization:
            suffix_static = self.prepare_suffix_static_context(state, curr_pad_masks, max_position_ids)
        while time >= -dt / 2:
            expanded_time = time.expand(batch_size)
            if suffix_static is None:
                v_t = self.denoise_step(
                    state,
                    curr_pad_masks,
                    past_key_values,
                    max_position_ids,
                    x_t.to(dtype),
                    expanded_time.to(dtype),
                )
            else:
                v_t = self.denoise_step_optimized(
                    suffix_static, past_key_values, x_t.to(dtype), expanded_time.to(dtype)
                )
            x_t = x_t + dt * v_t
            time += dt

        return x_t, None if not decode_image else middle_out

    def denoise_step(
        self,
        state: torch.Tensor,
        prefix_pad_masks: torch.Tensor,
        past_key_values: Any,
        max_prefix_position_ids: torch.Tensor,
        x_t: torch.Tensor,
        timestep: torch.Tensor,
    ) -> torch.Tensor:
        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, x_t, timestep)
        suffix_len = suffix_pad_masks.shape[1]
        batch_size = prefix_pad_masks.shape[0]
        prefix_len = prefix_pad_masks.shape[1]
        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
        position_ids = (
            torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device)
            .repeat(3, 1, 1)
            .to(max_prefix_position_ids)
            + max_prefix_position_ids
        )
        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
            full_att_2d_masks,
            dtype=suffix_embs.dtype,
        )
        outputs_embeds, _ = self.qwen3_vl_with_expert.forward(
            attention_mask=full_att_2d_masks_4d,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=[None, None, suffix_embs],
            use_cache=False,
        )
        suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32)
        return self.action_out_proj(suffix_out)

    def denoise_step_optimized(
        self,
        suffix_static: SuffixStaticContext,
        past_key_values: Any,
        x_t: torch.Tensor,
        timestep: torch.Tensor,
    ) -> torch.Tensor:
        state_emb = suffix_static.state_emb
        time_emb = create_sinusoidal_pos_embedding(
            timestep,
            self.action_in_proj.out_features,
            min_period=self.config.min_period,
            max_period=self.config.max_period,
            device=timestep.device,
        ).to(dtype=timestep.dtype)
        action_emb = self.action_in_proj(x_t)
        time_emb = time_emb[:, None, :].expand_as(action_emb)
        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
        action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb)))
        suffix_embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1)
        outputs_embeds, _ = self.qwen3_vl_with_expert.forward(
            attention_mask=suffix_static.full_att_2d_masks_4d,
            position_ids=suffix_static.position_ids,
            past_key_values=past_key_values,
            inputs_embeds=[None, None, suffix_embs],
            use_cache=False,
        )
        suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32)
        return self.action_out_proj(suffix_out)


class InternVLAA1Policy(nn.Module):
    _AUTO_WEIGHTS_IGNORE_PREFIXES = ["model.cosmos."]

    def __init__(
        self,
        config: InternVLAA1Config,
        *,
        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
    ) -> None:
        super().__init__()
        self.config = config
        self.input_builder = Qwen3VLInputBuilder(
            processor_model_name=processor_model_name,
            max_length=config.tokenizer_max_length,
        )
        cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths()
        self.model = InternVLAA1(
            config,
            cosmos_encoder_path=cosmos_encoder_path,
            cosmos_decoder_path=cosmos_decoder_path,
        )
        self.model.to(config.device)

    @classmethod
    def from_pretrained(
        cls,
        checkpoint_dir: str | Path,
        *,
        config: InternVLAA1Config | None = None,
        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
        strict: bool = False,
    ) -> InternVLAA1Policy:
        checkpoint_dir = Path(checkpoint_dir)
        if config is None:
            config = InternVLAA1Config.from_pretrained(checkpoint_dir)
        instance = cls(config, processor_model_name=processor_model_name)
        checkpoint_path = checkpoint_dir / "model.safetensors"
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"InternVLAA1 checkpoint not found: {checkpoint_path}")

        loader = AutoWeightsLoader(
            instance,
            ignore_unexpected_prefixes=cls._AUTO_WEIGHTS_IGNORE_PREFIXES,
        )
        with safe_open(str(checkpoint_path), framework="pt", device=str(config.device)) as f:
            file_keys = list(f.keys())
            loaded = loader.load_weights((name, f.get_tensor(name)) for name in file_keys)

        if strict:
            expected = {
                name
                for name, _ in instance.state_dict().items()
                if not any(name.startswith(prefix) for prefix in cls._AUTO_WEIGHTS_IGNORE_PREFIXES)
            }
            missing = sorted(expected - loaded)
            if missing:
                preview = ", ".join(missing[:10])
                suffix = " ..." if len(missing) > 10 else ""
                raise RuntimeError(f"Missing weights after AutoWeightsLoader load: {preview}{suffix}")

        instance.eval()
        return instance

    def to(self, *args, **kwargs):  # type: ignore[override]
        super().to(*args, **kwargs)
        self.model.cosmos.to(torch.bfloat16)
        self.model.action_out_proj.to(torch.float32)
        return self

    def _prepare_resized_histories(self, batch: dict[str, Any]) -> list[torch.Tensor]:
        resized_images: list[torch.Tensor] = []
        for i in range(3):
            image_history = batch[f"{OBS_IMAGES}.image{i}"]
            if image_history.shape[0] != 1:
                raise ValueError(f"InternVLAA1Policy only supports bs=1, got image batch {tuple(image_history.shape)}")
            resized_images.append(resize_with_pad(image_history[0], self.config.image_resolution))
        return resized_images

    def _preprocess_images(
        self,
        batch: dict[str, torch.Tensor],
        *,
        resized_histories: list[torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if resized_histories is None:
            resized_histories = self._prepare_resized_histories(batch)
        images = torch.stack(resized_histories, dim=0).unsqueeze(0)
        img_masks = torch.stack([batch[f"{OBS_IMAGES}.image{i}_mask"] for i in range(3)], dim=1)
        return images, img_masks

    def prepare_state(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        return pad_vector(batch[OBS_STATE], self.config.max_state_dim)

    def _get_task(self, batch: dict[str, Any]) -> str:
        task = batch[OBS_TASK]
        if isinstance(task, str):
            return task
        if isinstance(task, (list, tuple)):
            if len(task) != 1:
                raise ValueError(f"InternVLAA1Policy only supports bs=1, got {len(task)} tasks")
            item = task[0]
            if not isinstance(item, str):
                raise TypeError(f"Expected task string, got {type(item)!r}")
            return item
        raise TypeError(f"Unsupported task payload type: {type(task)!r}")

    def _prepare_qwen_prefix_inputs(
        self,
        resized_histories: list[torch.Tensor],
        task: str,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        qwen_inputs = self.input_builder.build(resized_histories, task)
        device = self.model.action_in_proj.weight.device
        return (
            qwen_inputs[f"{OBS_PREFIX}pixel_values"]
            .unsqueeze(0)
            .to(device=device, dtype=self.model.action_in_proj.weight.dtype),
            qwen_inputs[f"{OBS_PREFIX}image_grid_thw"].unsqueeze(0).to(device=device),
            qwen_inputs[f"{OBS_PREFIX}input_ids"].unsqueeze(0).to(device=device),
            qwen_inputs[f"{OBS_PREFIX}attention_mask"].unsqueeze(0).to(device=device),
        )

    @torch.inference_mode()
    def forward(
        self,
        batch: dict[str, Any],
        *,
        noise: torch.Tensor | None = None,
        decode_image: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        task = self._get_task(batch)
        resized_histories = self._prepare_resized_histories(batch)
        pixel_values, image_grid_thw, lang_tokens, lang_masks = self._prepare_qwen_prefix_inputs(
            resized_histories,
            task,
        )
        state = self.prepare_state(batch)
        images, img_masks = self._preprocess_images(batch, resized_histories=resized_histories)
        return self.model.sample_actions(
            images,
            img_masks,
            pixel_values,
            image_grid_thw,
            lang_tokens,
            lang_masks,
            state,
            noise=noise,
            decode_image=decode_image,
        )