src/inference_server/models/pi0fast_inference.py

import logging

import numpy as np
import torch
from lerobot.common.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
from lerobot.common.utils.utils import init_logging
from torchvision import transforms

from .base_inference import BaseInferenceEngine

logger = logging.getLogger(__name__)


class Pi0FastInferenceEngine(BaseInferenceEngine):
    """
    Pi0Fast (Physical Intelligence Fast) inference engine.

    Handles image preprocessing, joint normalization, and fast action prediction
    for Pi0Fast models with language instruction support.
    """

    def __init__(
        self,
        policy_path: str,
        camera_names: list[str],
        device: str | None = None,
        language_instruction: str | None = None,
    ):
        super().__init__(policy_path, camera_names, device)

        # Pi0Fast-specific configuration
        self.language_instruction = language_instruction
        self.supports_language = True

    async def load_policy(self):
        """Load the Pi0Fast policy from the specified path."""
        logger.info(f"Loading Pi0Fast policy from: {self.policy_path}")

        # Initialize hydra config for LeRobot
        init_logging()

        # Load the Pi0Fast policy
        self.policy = PI0FASTPolicy.from_pretrained(self.policy_path)
        self.policy.to(self.device)
        self.policy.eval()

        # Set up image transforms based on policy config
        if hasattr(self.policy, "config"):
            self._setup_image_transforms()

        self.is_loaded = True
        logger.info(f"Pi0Fast policy loaded successfully on {self.device}")

    def _setup_image_transforms(self):
        """Set up image transforms based on the policy configuration."""
        # Get image size from policy config
        config = self.policy.config
        image_size = getattr(config, "image_size", 224)

        # Create transforms for each camera
        for camera_name in self.camera_names:
            # Use policy-specific transforms if available
            if hasattr(self.policy, "image_processor"):
                self.image_transforms[camera_name] = self.policy.image_processor
            else:
                # Fall back to default transform
                self.image_transforms[camera_name] = transforms.Compose([
                    transforms.Resize((image_size, image_size)),
                    transforms.ToTensor(),
                    transforms.Normalize(
                        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                    ),
                ])

    async def predict(
        self, images: dict[str, np.ndarray], joint_positions: np.ndarray, **kwargs
    ) -> np.ndarray:
        """
        Run Pi0Fast inference to predict actions.

        Args:
            images: Dictionary of {camera_name: rgb_image_array}
            joint_positions: Current joint positions in LeRobot standard order
            task: Optional language instruction (overrides instance language_instruction)

        Returns:
            Array of predicted actions

        """
        if not self.is_loaded:
            msg = "Policy not loaded. Call load_policy() first."
            raise RuntimeError(msg)

        # Preprocess inputs
        processed_images = self.preprocess_images(images)
        processed_joints = self.preprocess_joint_positions(joint_positions)

        # Get language instruction
        task = kwargs.get("task", self.language_instruction)

        # Prepare batch inputs for Pi0Fast
        batch = self._prepare_batch(processed_images, processed_joints, task)

        # Run inference
        with torch.no_grad():
            action = self.policy.predict(batch)

            # Convert to numpy
            if isinstance(action, torch.Tensor):
                action = action.cpu().numpy()

            return action

    def _prepare_batch(
        self,
        images: dict[str, torch.Tensor],
        joints: torch.Tensor,
        task: str | None = None,
    ) -> dict:
        """
        Prepare batch inputs for Pi0Fast model.

        Args:
            images: Preprocessed images
            joints: Preprocessed joint positions
            task: Language instruction

        Returns:
            Batch dictionary for Pi0Fast model

        """
        batch = {}

        # Add images to batch
        for camera_name, image_tensor in images.items():
            # Add batch dimension if needed
            if len(image_tensor.shape) == 3:
                image_tensor = image_tensor.unsqueeze(0)
            batch[f"observation.images.{camera_name}"] = image_tensor

        # Add joint positions
        if len(joints.shape) == 1:
            joints = joints.unsqueeze(0)
        batch["observation.state"] = joints

        # Add language instruction if provided
        if task:
            batch["task"] = task

        return batch