VLA/wall_x/serving/policy/utils.py

from typing import Dict, List
import logging
import numpy as np
from wall_x.data.utils import preprocesser_call
from qwen_vl_utils.vision_process import smart_resize
import torch
from PIL import Image
from transformers import BatchFeature

logger = logging.getLogger(__name__)


def prepare_batch(
    obs: Dict,
    processor,
    normalizer_propri,
    camera_key: List[str],
    agent_pos_dim,
    action_dim,
    pred_horizon,
    fixed_action_dim,
    max_length,
    image_factor: int,
    min_pixels: int,
    max_pixels: int,
    predict_mode: str = "fast",
    device: str = "cuda",
) -> BatchFeature:
    """Prepare observation into model input format.

    Args:
        obs: Dictionary containing:
            - 'camera_key_0' : image 0
            - 'camera_key_1' : image 1
            ...
            - 'prompt': Text prompt
            - 'state': Robot state/proprioception
            - 'dataset_names': Dataset names

    Returns:
        BatchFeature object ready for model input
    """
    # Handle images - can be single image, list of images, or dict of images
    images = []
    images = [obs[key] for key in camera_key]
    # Convert numpy arrays to PIL Images
    processed_images = []
    for img in images:
        if isinstance(img, np.ndarray):
            # Debug: Log the shape and dtype
            logger.debug(f"Image shape: {img.shape}, dtype: {img.dtype}")

            # Handle unexpected dimensions - squeeze if needed
            if img.ndim > 3:
                logger.warning(
                    f"Image has {img.ndim} dimensions, squeezing extra dimensions"
                )
                img = np.squeeze(img)

            # Verify shape is valid for PIL
            if img.ndim == 2:
                # Grayscale image
                pass
            elif img.ndim == 3:
                # Check if channel dimension is first or last
                if img.shape[0] == 3 or img.shape[0] == 1:
                    # Channels first, transpose to channels last
                    img = np.transpose(img, (1, 2, 0))
                elif img.shape[2] == 3 or img.shape[2] == 1:
                    # Already channels last
                    pass
                else:
                    raise ValueError(
                        f"Unexpected image shape: {img.shape}. Expected (H, W, C) or (C, H, W)"
                    )
            else:
                raise ValueError(
                    f"Invalid image dimensions: {img.ndim}. Expected 2 or 3 dimensions, got shape {img.shape}"
                )

            # Convert to PIL Image
            if img.dtype == np.uint8:
                img = Image.fromarray(img)
            else:
                img = Image.fromarray((img * 255).astype(np.uint8))
        processed_images.append(img)

    # print("processed_images:",processed_images)
    # Apply smart resize to images
    resized_images = process_images(
        processed_images, image_factor, min_pixels, max_pixels
    )

    # Handle text prompt - format with vision tokens
    instruction = obs["prompt"]
    formatted_text = format_text_with_vision_tokens(
        instruction, camera_key, predict_mode, pred_horizon
    )

    # Use processor to prepare inputs
    inputs = preprocesser_call(
        processor=processor,
        text=[formatted_text],
        images=[resized_images],
        videos=None,
        padding=True,
        truncation=True,
        return_tensors="pt",
        max_length=max_length,
    )

    action_token_id = processor.tokenizer.convert_tokens_to_ids("<|action|>")
    moe_token_types = inputs.input_ids == action_token_id
    inputs["moe_token_types"] = torch.tensor(moe_token_types)

    # obs["dataset_names"]="libero_all"

    # Handle robot state/proprioception if available
    if "state" in obs:
        state = obs["state"]
        if isinstance(state, np.ndarray):
            state = torch.from_numpy(state).float()
        elif not isinstance(state, torch.Tensor):
            state = torch.tensor(state, dtype=torch.float32)

        # Add batch dimension if needed
        if state.dim() == 1:
            state = state.unsqueeze(0)
        if state.dim() == 2:
            state = state.unsqueeze(1)  # [batch, 1, state_dim]

        # Pad to 20 dimensions if needed (same as training)
        # if state.shape[-1] < 20:
        #     padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1])
        #     state = torch.cat([state, padding], dim=-1)

        # Create mask for valid dimensions
        agent_pos_mask = torch.ones_like(state)
        if state.shape[-1] > agent_pos_dim:
            agent_pos_mask[:, :, agent_pos_dim:] = 0

        normalizer_propri.normalize_data(state, [obs["dataset_names"]] * state.shape[0])

        inputs["proprioception"] = state
        inputs["agent_pos_mask"] = agent_pos_mask

    # Add dataset name (required by model)
    inputs["dataset_names"] = [obs["dataset_names"]] * state.shape[0]

    # Move all tensors to device
    for key in inputs:
        if isinstance(inputs[key], torch.Tensor):
            inputs[key] = inputs[key].to(device)

    dof_mask = torch.ones([state.shape[0], pred_horizon, fixed_action_dim])
    dof_mask[:, :, action_dim:] = 0

    inputs["dof_mask"] = dof_mask

    # Convert to BatchFeature to maintain consistency with training pipeline
    return BatchFeature(data=dict(inputs)).to(device)


def process_images(
    images: List[Image.Image], image_factor: int, min_pixels: int, max_pixels: int
) -> List[Image.Image]:
    """Process images with smart resize following the data loading pattern.

    Args:
        images: List of PIL Images

    Returns:
        List of resized PIL Images
    """
    resized_images = []
    for img_pil in images:

        orig_width, orig_height = img_pil.size
        target_size = 256
        if target_size != -1:
            # Maintain aspect ratio logic
            if orig_width > orig_height:  # Landscape image
                new_width = target_size
                new_height = int(target_size * orig_height / orig_width)
            else:  # Portrait image
                new_height = target_size
                new_width = int(target_size * orig_width / orig_height)
            img_pil = img_pil.resize((new_width, new_height))

        # Apply smart scaling (Qwen logic)
        current_width, current_height = img_pil.size
        resized_height, resized_width = smart_resize(
            current_height,
            current_width,
            factor=image_factor,
            min_pixels=min_pixels,
            max_pixels=max_pixels,
        )

        resized_img = img_pil.resize((resized_width, resized_height))
        resized_images.append(resized_img)

    return resized_images


def format_text_with_vision_tokens(
    instruction: str,
    camera_key: List[str],
    predict_mode: str = "diffusion",
    pred_horizon: int = 32,
) -> str:
    """Format text prompt with vision tokens for the model.

    Args:
        instruction: Task instruction text
        camera_key: List of camera names

    Returns:
        Formatted text with special tokens
    """
    # Special tokens for formatting
    role_start_symbol = "<|im_start|>"
    role_end_symbol = "<|im_end|>"
    vision_start_symbol = "<|vision_start|>"
    vision_end_symbol = "<|vision_end|>"
    image_pad_symbol = "<|image_pad|>"
    propri_symbol = "<|propri|>"
    action_symbol = "<|action|>"
    action_fast_symbol = "<|action_fast|>"

    # Camera name mapping
    camera_name_mapping = {
        "front_view": "front view",
        "face_view": "front view",
        "left_wrist_view": "left wrist view",
        "right_wrist_view": "right wrist view",
        "top_view": "top view",
        "wall_view": "wall view",
    }

    # System prologue
    prologue = (
        f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n"
    )

    # User request with observation
    user_request = f"{role_start_symbol}user\nObservation:"
    if camera_key:
        for cam_name in camera_key:
            view_name = camera_name_mapping.get(cam_name, cam_name)
            user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}"
    user_request += "\nInstruction:"

    text_prompt = (
        f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n"
    )
    user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
    assistant_output = (
        f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n"
    )
    if predict_mode == "diffusion":
        assistant_output = f"{role_start_symbol}assistant\n{action_symbol * pred_horizon}{role_end_symbol}\n"
    complete_text = prologue + user_message + assistant_output

    return complete_text