code/model/framework/M1.py

# Copyright 2025 InternVLA-M1. All rights reserved.
# Modified by [Jinhui YE/ HKUST University] in [2025]. 
# Modification: [add fake sample and predict_action to match with starVLA].
"""
InternVLA M1 framework:
Vision-Language-Action diffusion model integrating:
  - Qwen2.5 vision-language backbone
  - Layer-wise QFormer aggregation
  - DINO multi-view visual encoder
  - DiT diffusion head for future action sequence prediction
Primary goal: predict continuous future actions conditioned on multi-view images + instruction.
"""

from typing import List
from tqdm import tqdm
from typing import List, Optional, Tuple
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
from qwen_vl_utils import process_vision_info


from starVLA.training.trainer_utils import initialize_overwatch
from starVLA.model.tools import FRAMEWORK_REGISTRY


logger = initialize_overwatch(__name__)

# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
IGNORE_INDEX = -100

from starVLA.model.framework.base_framework import baseframework
from starVLA.model.modules.vlm import get_vlm_model
from starVLA.model.modules.projector.QFormer import get_layerwise_qformer
from starVLA.model.modules.action_model.DiTActionHeader import get_action_model
from starVLA.model.modules.dino_model.dino import get_dino_model
from starVLA.training.trainer_utils.trainer_tools import resize_images


@FRAMEWORK_REGISTRY.register("InternVLA-M1")
class InternVLA_M1(baseframework):
    """
    Multimodal vision-language-action model.

    Components:
      - Qwen2.5 VL interface for fused language/vision token embeddings
      - Layer-wise QFormer for multi-layer feature aggregation
      - DINO encoder for dense multi-view spatial tokens
      - DiT diffusion head for future action sequence modeling

    Focus: Predict future continuous actions conditioned on images + instruction.
    """

    def __init__(
        self,
        config: Optional[dict] = None,
        **kwargs,
    ) -> None:
        """
        Construct all submodules and cache key configuration values.

        Args:
            config: Hierarchical configuration (OmegaConf/dict) containing framework + trainer sections.
            **kwargs: Reserved for future overrides (unused).
        """
        super().__init__()
        self.config = config
        self.qwen_vl_interface = get_vlm_model(config=self.config)
        self.layer_qformer = get_layerwise_qformer(config=self.config)
        self.action_model = get_action_model(config=self.config)
        self.dino_encoder = get_dino_model(
            backone_name=getattr(self.config.framework.dino, "dino_backbone", "dinov2_vits14")
        )
        self.dino_pro = nn.Linear(
            in_features=self.dino_encoder.num_channels, out_features=self.qwen_vl_interface.model.config.hidden_size
        )

        self.future_action_window_size = config.framework.action_model.future_action_window_size
        self.past_action_window_size = config.framework.action_model.past_action_window_size

    def forward(
        self,
        examples: List[dict] = None,
        **kwargs,
    ) -> Tuple:
        """
        Forward pass for training (diffusion objective).

        Flow:
          1. Build QwenVL inputs (images + instruction tokens)
          2. Extract hidden states from configured layer range
          3. Encode images with DINO, flatten multi-view tokens and project
          4. Concatenate per-layer language tokens with visual tokens
          5. Fuse via layer-wise QFormer -> action condition embeddings
          6. Prepare repeated future action windows (for diffusion efficiency)
          7. Predict noise and compute diffusion loss

        Args:
            examples: List[dict], each dict requires:
                - image: List[PIL.Image] (multi-view)
                - lang: str instruction
                - action: np.ndarray or list shaped [T, action_dim]
            **kwargs: Reserved.

        Returns:
            dict:
                action_loss (torch.Tensor): Scalar diffusion noise prediction loss.
        """
        batch_images = [example["image"] for example in examples]  #  [B，[PLT]]
        instructions = [example["lang"] for example in examples]  # [B, str]
        actions = [example["action"] for example in examples]  # label [B， len, 7]

        # Step 1: QWenVL input format
        qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            qwenvl_outputs = self.qwen_vl_interface(
                **qwen_inputs,
                output_attentions=False,
                output_hidden_states=True,
                return_dict=True,
            )
            pass

        # Step 2: DINO Forward
        image_tensors = self.dino_encoder.prepare_dino_input(batch_images)  #
        B = len(batch_images)
        dino_features = self.dino_encoder(image_tensors)  # DINO output is [B*num_view, token, dim]
        dino_encoded_features = dino_features.reshape(B, -1, dino_features.shape[-1])  # [B, num_view * token, dim]
        dino_encoded_features = self.dino_pro(dino_encoded_features)  # [B, num_view * token, hidden_size]

        # Step 3: aggregation condition for Action expert
        start_layer = self.config.framework.layer_qformer.qformer_start_layer
        end_layer = self.config.framework.layer_qformer.qformer_end_layer
        condition_features = qwenvl_outputs.hidden_states[start_layer:end_layer]

        cat_conditions = []
        for layer_index in range(len(condition_features)):
            layer_features = condition_features[layer_index]  # [B, n_qformer_token, D]
            layer_features = torch.cat(
                [layer_features, dino_encoded_features], dim=1
            )  # [B, n_qformer_token + num_view * token, D]
            cat_conditions.append(layer_features)

        action_condition = self.layer_qformer(cat_conditions)  # [B, 64, D_action]

        # Step 4: Action Expert Forward and Loss
        with torch.autocast("cuda", dtype=torch.float32):

            # here is a tips to accelerate training speed, by repeating each sample for several times @ref to CogACT
            actions = torch.tensor(np.array(actions), device=action_condition.device)  # [B, chunk, 7]
            actions_future = actions[:, -(self.future_action_window_size + 1) :, :]

            # tips: Repeat 'actions' 'repeated_diffusion_steps' times, resulting in [repeated_diffusion_steps*B, T, D]
            repeated_diffusion_steps = (
                self.config.trainer.get("repeated_diffusion_steps", 4) if self.config and self.config.trainer else 4
            )
            actions_repeated = actions_future.repeat(repeated_diffusion_steps, 1, 1)
            action_condition = action_condition.repeat(
                repeated_diffusion_steps, 1, 1
            )  # [repeated_diffusion_steps*B, T, D_action]

            # DiT noise add and predict
            noise_pred, noise, timestep = self.action_model(actions_repeated, action_condition)

            # perdition loss
            action_loss = self.action_model.loss(noise_pred, noise)

        return {"action_loss": action_loss}

    @torch.inference_mode()
    def predict_action(
        self,
        batch_images: List[List[Image.Image]],  # B * List of PIL Image as [view1, view2]
        instructions: List[str],
        cfg_scale: float = 1.5,
        use_ddim: bool = True,
        num_ddim_steps: int = 5,
        resize_image = [224, 224],
        **kwargs: str,
    ) -> np.ndarray:
        """
        Inference: generate future normalized action sequence via diffusion sampling.

        Steps:
          1. Resize images to training resolution (if specified)
          2. Encode with QwenVL (hidden states retained)
          3. Extract DINO tokens and project to vlm hidden size
          4. Build multi-layer fused QwenVL and DINO features via QFormer
          5. Run diffusion sampling (DDIM optional, CFG optional)
          6. Return normalized action trajectory

        Args:
            batch_images: List of samples; each sample is List[PIL.Image] (multi-view).
            instructions: List[str] natural language task instructions.
            cfg_scale: >1 enables classifier-free guidance (scales conditional vs unconditional).
            use_ddim: Whether to use DDIM deterministic sampling.
            num_ddim_steps: Number of DDIM steps if enabled.
            **kwargs: Reserved.

        Returns:
            dict:
                normalized_actions (np.ndarray): Shape [B, T, action_dim], diffusion-sampled normalized actions.
        """
        # align obs and lang # is policy's duty to make sure the image size?
        train_obs_image_size = getattr(self.config.datasets.vla_data, "image_size", None)
        if train_obs_image_size:
            batch_images = resize_images(batch_images, target_size=train_obs_image_size)
        instructions = [instruction.lower() for instruction in instructions]

        inferface_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
        qwen_inputs = inferface_inputs

        with torch.autocast("cuda", dtype=torch.bfloat16):
            qwenvl_outputs = self.qwen_vl_interface(
                **qwen_inputs,
                output_hidden_states=True,
                return_dict=True,
            )

            B = len(batch_images) # dino don't have smart resize in processing
            image_tensors = self.dino_encoder.prepare_dino_input(batch_images)
            dino_features = self.dino_encoder(image_tensors)

            B = len(batch_images)
            dino_encoded_features = dino_features.reshape(B, -1, dino_features.shape[-1])  # [B, num_view * token, dim]
            dino_encoded_features = self.dino_pro(dino_encoded_features)  # [B, 256, D]

        with torch.autocast("cuda", dtype=torch.bfloat16):

            start_layer = self.config.framework.layer_qformer.qformer_start_layer
            end_layer = self.config.framework.layer_qformer.qformer_end_layer
            condition_features = qwenvl_outputs.hidden_states[start_layer:end_layer]
            cat_conditions = []
            for layer_index in range(len(condition_features)):
                layer_features = condition_features[layer_index]  # [B, n_qformer_token, D]
                layer_features = torch.cat(
                    [layer_features, dino_encoded_features], dim=1
                )  # [B, n_qformer_token + num_view * token, D]
                cat_conditions.append(layer_features)

            action_condition_feature = self.layer_qformer(cat_conditions)  # [B, 64, D_action]

            using_cfg = cfg_scale > 1.0

            model_dtype = next(self.action_model.net.parameters()).dtype
            B = action_condition_feature.shape[0]

            # Sample random noise
            noise = torch.randn(
                B,
                self.future_action_window_size + 1,
                self.action_model.in_channels,
                device=action_condition_feature.device,
            ).to(
                model_dtype
            )  # [B, T, D]

            # Setup classifier-free guidance:
            if using_cfg:
                noise = torch.cat([noise, noise], 0)  # [2,16,7]
                uncondition = self.action_model.net.z_embedder.uncondition  # [64, 768]
                uncondition_shape = uncondition.shape
                uncondition = uncondition.unsqueeze(0)  # [1, 64, D]
                uncondition = uncondition.expand(
                    B, uncondition_shape[0], uncondition_shape[1]
                )  # [B, n_qformer_token, D]
                z = torch.cat([action_condition_feature, uncondition], 0)  # [2, 64, 768]
                cfg_scale = cfg_scale
                model_kwargs = dict(z=z, cfg_scale=cfg_scale)
                sample_fn = self.action_model.net.forward_with_cfg
            else:
                model_kwargs = dict(z=action_condition_feature)
                sample_fn = self.action_model.net.forward

            # DDIM Sampling
            if use_ddim and num_ddim_steps is not None:
                if self.action_model.ddim_diffusion is None:
                    self.action_model.create_ddim(ddim_step=num_ddim_steps)
                samples = self.action_model.ddim_diffusion.ddim_sample_loop(
                    sample_fn,
                    noise.shape,
                    noise,
                    clip_denoised=False,
                    model_kwargs=model_kwargs,
                    progress=False,
                    device=action_condition_feature.device,
                    eta=0.0,
                )

            if using_cfg:
                samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
            normalized_actions = samples.cpu().numpy()
            
            raw_actions = None
     
        return {"normalized_actions": normalized_actions}  # [B, T, action_dim]


    @torch.inference_mode()
    def chat_with_M1(
        self,
        image: Image.Image,
        text: str,
        max_new_tokens: int = 128,
        device: Optional[str] = "cuda",
    ) -> List[str]:
        processor = getattr(self.qwen_vl_interface, "processor", None)
        model = getattr(self.qwen_vl_interface, "model", None)
        # if processor is None or model is None:
        #     raise RuntimeError("qwen_vl_interface 缺少 processor 或 model。")

        messages0 = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": image,
                    },
                    {"type": "text", "text": text},
                ],
            }
        ]

        messages = [messages0]
        # text info
        texts = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) for msg in messages]
        # visual info
        image_inputs, video_inputs = process_vision_info(messages)

        # tokenizer
        inputs = processor(
            text=texts,
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ).to(device)

        model.eval()
        generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)
        generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
        outputs = processor.batch_decode(
            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )
        return outputs

if __name__ == "__main__":
    from omegaconf import OmegaConf
    import debugpy
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--config_yaml", type=str, default="./starVLA/config/training/starvla_cotrain_oxe.yaml", help="Path to YAML config")
    args, clipargs = parser.parse_known_args()

    debugpy.listen(("0.0.0.0", 10092))
    print("🔍 Rank 0 waiting for debugger attach on port 10092...")
    debugpy.wait_for_client()

    cfg = OmegaConf.load(args.config_yaml)

    # try get model
    model = InternVLA_M1(cfg)
    print(model)


    # fake sample 
    image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
    # Create a sample
    sample = {
        "action": np.random.uniform(-1, 1, size=(16, 7)).astype(np.float16), # action_chunk, action_dim
        "image": [image, image], # two views
        "lang": "This is a fake instruction for testing.",
        # "state" : np.random.uniform(-1, 1, size=(1, 7)).astype(np.float16), # chunk, state_dim
    }

    batch  = [sample, sample]  # batch size 2
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    forward_output = model(batch)
    action_loss = forward_output['action_loss']
    print(f"Action Loss: {action_loss.item()}")

    # test predict action
    predict_output = model.predict_action(batch_images=[batch[0]["image"]], instructions=[batch[0]["lang"]])
    normalized_actions = predict_output['normalized_actions']
    print(f"Unnormalized Action: {normalized_actions}")




    # model_path = "./results/Checkpoints/1_need/0906_bestvla_retrain_sota2/checkpoints/steps_50000_pytorch_model.pt"
    # state_dict = torch.load(model_path, map_location="cpu")

    # model.load_state_dict(state_dict, strict=True)

    # # try forward model
    # # can be fake sample， but here get from dataloader for simpler
    # from starVLA.dataloader.lerobot_datasets import get_vla_dataset, collate_fn

    # vla_dataset_cfg = cfg.datasets.vla_data
    # dataset = get_vla_dataset(data_cfg=vla_dataset_cfg)

    # from torch.utils.data import DataLoader

    # train_dataloader = DataLoader(
    #     dataset,
    #     batch_size=2,
    #     num_workers=1,  # For Debug
    #     collate_fn=collate_fn,
    # )

    # for batch in tqdm(train_dataloader, desc="Processing Batches"):
    #     batch
    #     break

    # # try get model
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # model = model.to(device)
    # model(batch)