inference_human_prediction.py

import os
import sys
import cv2
import math
import json
import torch
import argparse
import numpy as np
from PIL import Image
from PIL import ImageOps
from pathlib import Path
import multiprocessing as mp
from vitra.models import VITRA_Paligemma, load_model
from vitra.utils.data_utils import resize_short_side_to_target, load_normalizer, recon_traj
from vitra.utils.config_utils import load_config
from vitra.datasets.human_dataset import pad_state_human, pad_action
from scipy.spatial.transform import Rotation as R
from vitra.datasets.dataset_utils import (
    compute_new_intrinsics_resize, 
    calculate_fov,
    ActionFeature,
    StateFeature,
)

repo_root = Path(__file__).parent.parent  # VITRA/
sys.path.insert(0, str(repo_root))

from visualization.visualize_core import HandVisualizer, normalize_camera_intrinsics, save_to_video, Renderer, process_single_hand_labels
from visualization.visualize_core import Config as HandConfig

def main():
    """
    Main execution function for hand action prediction and visualization.
    
    This function uses a multi-process architecture to separate hand reconstruction
    and VLA inference into independent processes, preventing CUDA conflicts.
    
    Workflow:
    1. Parse command-line arguments and load model configurations
    2. Initialize persistent services:
       - HandReconstructionService: Runs HAWOR + MOGE models in separate process
       - VLAInferenceService: Runs VLA model in separate process
    3. Load or reconstruct hand state:
       - Uses precomputed .npy file if available (same stem as image)
       - Otherwise runs hand reconstruction service
    4. Prepare input data:
       - Load and resize image
       - Extract hand state (translation, rotation, pose) for left/right hands
       - Create state and action masks based on which hands to predict
    5. Run VLA inference to predict future hand actions (multiple samples for diversity)
    6. Reconstruct absolute hand trajectories from relative actions
    7. Visualize predicted hand motions using MANO hand model
    8. Generate grid layout video showing all samples and save to file
    9. Cleanup: Shutdown persistent services and free GPU memory

    """
    parser = argparse.ArgumentParser(description="Hand VLA inference and visualization.")
    
    # Model Configuration
    parser.add_argument('--config_path', type=str, required=True, help='Path to model configuration JSON file')
    parser.add_argument('--model_path', type=str, default=None, help='Path to model checkpoint (overrides config)')
    parser.add_argument('--statistics_path', type=str, default=None, help='Path to normalization statistics JSON (overrides config)')
    
    # Input/Output
    parser.add_argument('--image_path', type=str, required=True, help='Path to input image file')
    parser.add_argument('--hand_path', type=str, default=None, help='Path to hand state .npy file (optional, will run reconstruction if not provided)')
    parser.add_argument('--video_path', type=str, default='./example_human_inf.mp4', help='Path to save output visualization video')
    
    # Hand Reconstruction Models
    parser.add_argument('--hawor_model_path', type=str, default='./weights/hawor/checkpoints/hawor.ckpt', help='Path to HAWOR model weights')
    parser.add_argument('--detector_path', type=str, default='./weights/hawor/external/detector.pt', help='Path to hand detector model')
    parser.add_argument('--moge_model_name', type=str, default='Ruicheng/moge-2-vitl', help='MOGE model name from Hugging Face')
    parser.add_argument('--mano_path', type=str, default='/home/t-qixiuli/repo/VITRA/weights/mano', help='Path to MANO model files')
    # parser.add_argument('--output_path', type=str, default='./recon_results.npy', help='Path to save reconstruction results')
    
    # Prediction Settings
    parser.add_argument('--use_left', action='store_true', help='Enable left hand prediction')
    parser.add_argument('--use_right', action='store_true', help='Enable right hand prediction')
    parser.add_argument('--instruction', type=str, default="Left hand: Put the trash into the garbage. Right hand: None.", help='Text instruction for hand motion')
    parser.add_argument('--sample_times', type=int, default=4, help='Number of action samples to generate for diversity')
    parser.add_argument('--fps', type=int, default=8, help='Frames per second for output video')
    
    # Advanced Options
    parser.add_argument('--save_state_local', action='store_true', help='Save hand state locally as .npy file')

    # === Environment Configuration ===
    # Disable tokenizers parallelism to avoid deadlocks in multi-process data loading
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    args = parser.parse_args()
    
    # Validate that at least one hand is selected
    if not args.use_left and not args.use_right:
        raise ValueError("At least one of --use_left or --use_right must be specified.")
    
    # Load configs
    configs = load_config(args.config_path)
    
    # Override config with command-line arguments if provided
    if args.model_path is not None:
        configs['model_load_path'] = args.model_path
    if args.statistics_path is not None:
        configs['statistics_path'] = args.statistics_path

    # Check if a precomputed hand reconstruction .npy (same stem as image) exists.
    image_path_obj = Path(args.image_path)
    npy_path = image_path_obj.with_suffix('.npy')

    # Initialize services
    print("Initializing services...")
    if npy_path.exists():
        # If precomputed .npy exists, load hand state from it
        # Hand State File Format
        # The hand state is stored as a .npy file containing a Python dictionary with the following structure:
        # Format: .npy
        # Content: dictionary with MANO-based hand pose parameters and camera FOV.
        # Structure:
        # {
        #     'left': {
        #         0: {
        #             'hand_pose': np.ndarray,      # [15, 3, 3] rotation matrices for MANO joints
        #             'global_orient': np.ndarray,  # [3, 3] global rotation matrix
        #             'transl': np.ndarray,         # [3] root translation in camera coordinates
        #             'beta': np.ndarray            # [10] MANO shape parameters
        #         }
        #     },
        #     'right': {                           # Same structure as 'left'
        #         0: {
        #             ...
        #         }
        #     },
        #     'fov_x': float                       # Horizontal field of view (in degrees)
        # }

        print(f"Found precomputed hand state results: {npy_path}. Using the state instead of running hand recon.")
        hand_data = np.load(npy_path, allow_pickle=True).item()

        hand_recon_service = None
    else:
        print(f"No precomputed hand state .npy found at {npy_path}. Starting hand reconstruction service.")
        
        # Start hand reconstruction service
        hand_recon_service = HandReconstructionService(args)
        hand_data = None


    # Start VLA service (normalizer and model are loaded inside the service)
    vla_service = VLAInferenceService(configs)
    
    # Visualization setup
    hand_config = HandConfig(args)
    hand_config.FPS = args.fps
    visualizer = HandVisualizer(hand_config, render_gradual_traj=False)

    try:
        if hand_data is None:
            # Run hand reconstruction using service
            print("Running hand reconstruction...")
            hand_data = hand_recon_service.reconstruct(args.image_path)
            if args.save_state_local:
                # Save hand state locally as .npy
                np.save(npy_path, hand_data, allow_pickle=True)
                print(f"Saved reconstructed hand state to {npy_path}")

        # Load and process image
        image = Image.open(args.image_path)
        ori_w, ori_h = image.size

        # Handle EXIF orientation if present (fixes image rotation issues)
        try:
            image = ImageOps.exif_transpose(image)
        except Exception:
            pass  # If EXIF handling fails, continue with original image

        image_resized = resize_short_side_to_target(image, target=224)
        w, h = image_resized.size

        use_right = args.use_right
        use_left = args.use_left

        # Initialize state
        current_state_left = None
        current_state_right = None
        
        if use_right:
            current_state_right, beta_right, fov_x, _ = get_state(hand_data, hand_side='right')
        if use_left:
            current_state_left, beta_left, fov_x, _ = get_state(hand_data, hand_side='left')
        
        fov_x = fov_x * np.pi /180
        f_ori = ori_w / np.tan(fov_x / 2) /2
        fov_y = 2 * np.arctan(ori_h / (2 * f_ori))

        f = w / np.tan(fov_x / 2) /2
        intrinsics = np.array([
            [f, 0, w/2],
            [0, f, h/2],
            [0, 0, 1]
        ])

        # Concatenate left and right hand states, filling with zeros if one is None
        if current_state_left is None and current_state_right is None:
            raise ValueError("Both current_state_left and current_state_right are None")
        
        state_left = current_state_left if use_left else np.zeros_like(current_state_right)
        beta_left = beta_left if use_left else np.zeros_like(beta_right)
        state_right = current_state_right if use_right else np.zeros_like(current_state_left)
        beta_right = beta_right if use_right else np.zeros_like(beta_left)
        
        state = np.concatenate([state_left, beta_left, state_right, beta_right], axis=0)
        state_mask = np.array([use_left, use_right], dtype=bool)
        
        # Note: chunk_size needs to be determined from config
        chunk_size = configs.get('fwd_pred_next_n', 16)  # Default to 16 if not in config
        action_mask = np.tile(np.array([[use_left, use_right]], dtype=bool), (chunk_size, 1)) 

        fov = np.array([fov_x, fov_y], dtype=np.float32)
        # Convert resized image to numpy for VLA inference
        image_resized_np = np.array(image_resized)

        # Use instruction from command-line argument
        instruction = args.instruction

        # Model inference using service (includes normalization, padding, and unnormalization)
        print(f"Running VLA inference...")
        sample_times = args.sample_times
        unnorm_action = vla_service.predict(
            image=image_resized_np,
            instruction=instruction,
            state=state,
            state_mask=state_mask,
            action_mask=action_mask,
            fov=fov,
            num_ddim_steps=10,
            cfg_scale=5.0,
            sample_times=sample_times,
        )
        
        fx_exo = intrinsics[0, 0]
        fy_exo = intrinsics[1, 1]
        renderer = Renderer(w, h, (fx_exo, fy_exo), 'cuda')

        T = len(action_mask) + 1
        traj_right_list = np.zeros((sample_times, T, 51), dtype=np.float32)
        traj_left_list = np.zeros((sample_times, T, 51), dtype=np.float32)

        traj_mask = np.tile(np.array([[use_left, use_right]], dtype=bool), (T, 1)) 
        left_hand_mask = traj_mask[:, 0]
        right_hand_mask = traj_mask[:, 1]

        # Masks indicating which hands are used in the trajectory
        hand_mask = (left_hand_mask, right_hand_mask)

        all_rendered_frames = []
        
        # Reconstruct trajectories and visualize for each sample
        for i in range(sample_times):
            traj_right = traj_right_list[i]
            traj_left = traj_left_list[i]
            
            if use_left:
                traj_left = recon_traj(
                    state=state_left,
                    rel_action=unnorm_action[i, :, 0:51],
                )
            if use_right:
                traj_right = recon_traj(
                    state=state_right,
                    rel_action=unnorm_action[i, :, 51:102],
                )
        
            left_hand_labels = {
                'transl_worldspace': traj_left[:, 0:3],
                'global_orient_worldspace': R.from_euler('xyz', traj_left[:, 3:6]).as_matrix(),
                'hand_pose': euler_traj_to_rotmat_traj(traj_left[:, 6:51], T),
                'beta': beta_left,
                }
            right_hand_labels = {
                'transl_worldspace': traj_right[:, 0:3],
                'global_orient_worldspace': R.from_euler('xyz', traj_right[:, 3:6]).as_matrix(),
                'hand_pose': euler_traj_to_rotmat_traj(traj_right[:, 6:51], T),
                'beta': beta_right,
            }

            verts_left_worldspace, _ = process_single_hand_labels(left_hand_labels, left_hand_mask, visualizer.mano, is_left=True)
            verts_right_worldspace, _ = process_single_hand_labels(right_hand_labels, right_hand_mask, visualizer.mano, is_left=False)

            hand_traj_wordspace = (verts_left_worldspace, verts_right_worldspace)
            
            R_w2c = np.broadcast_to(np.eye(3), (T, 3, 3)).copy()
            t_w2c = np.zeros((T, 3, 1), dtype=np.float32)

            extrinsics = (R_w2c, t_w2c)

            # Use resized image for visualization (convert RGB to BGR)
            image_bgr = image_resized_np[..., ::-1]
            resize_video_frames = [image_bgr] * T
            save_frames = visualizer._render_hand_trajectory(
                resize_video_frames,
                hand_traj_wordspace,
                hand_mask,
                extrinsics,
                renderer,
                mode='first'
            )
        
            all_rendered_frames.append(save_frames)
        
        # Concatenate all samples spatially into a single video
        # all_rendered_frames: list of sample_times frame lists
        # Each frame list has T frames
        num_frames = len(all_rendered_frames[0])
        
        # Determine grid layout (e.g., 2x2 for 4 samples)
        grid_cols = math.ceil(math.sqrt(sample_times))
        grid_rows = math.ceil(sample_times / grid_cols)
        
        # Combine different samples in one video
        combined_frames = []
        for frame_idx in range(num_frames):
            # Collect all sample frames at this time step
            sample_frames = [all_rendered_frames[i][frame_idx] for i in range(sample_times)]
            
            # Pad with black frames if needed to fill the grid
            while len(sample_frames) < grid_rows * grid_cols:
                black_frame = np.zeros_like(sample_frames[0])
                sample_frames.append(black_frame)
            
            # Arrange frames in grid
            rows = []
            for row_idx in range(grid_rows):
                row_frames = sample_frames[row_idx * grid_cols:(row_idx + 1) * grid_cols]
                row_concat = np.concatenate(row_frames, axis=1)  # Concatenate horizontally
                rows.append(row_concat)
            
            # Concatenate rows vertically
            combined_frame = np.concatenate(rows, axis=0)
            combined_frames.append(combined_frame)

        # Save combined video
        save_to_video(combined_frames, f'{args.video_path}', fps=hand_config.FPS)
        print(f"Combined video with {sample_times} samples saved to {args.video_path}")
    
    finally:
        # Cleanup persistent services
        print("Shutting down services...")
        if hand_recon_service is not None:
            hand_recon_service.shutdown()
        vla_service.shutdown()
        print("All services shut down successfully")
    

def get_state(hand_data, hand_side='right'):
    """
    Load and extract hand state from hand data.
    
    Args:
        hand_data (dict): Dictionary containing hand data
        hand_side (str): Which hand to extract, either 'left' or 'right'. Default is 'right'.
        
    Returns:
        tuple: (state_t0, beta, fov_x, None) where:
            - state_t0 (np.ndarray): Hand state [51] containing translation (3), 
                                     global rotation (3 euler angles), and hand pose (45 euler angles)
            - beta (np.ndarray): MANO shape parameters [10]
            - fov_x (float): Horizontal field of view in degrees
            - None: Placeholder for optional text annotations
    """
    if hand_side not in ['left', 'right']:
        raise ValueError(f"hand_side must be 'left' or 'right', got '{hand_side}'")
    
    hand_pose_t0 = hand_data[hand_side][0]['hand_pose']
    hand_pose_t0_euler = R.from_matrix(hand_pose_t0).as_euler('xyz', degrees=False) # [15, 3]
    hand_pose_t0_euler = hand_pose_t0_euler.reshape(-1)  # [45]
    global_orient_mat_t0 = hand_data[hand_side][0]['global_orient']
    R_t0_euler = R.from_matrix(global_orient_mat_t0).as_euler('xyz', degrees=False)  # [3]
    transl_t0 = hand_data[hand_side][0]['transl']  # [3]
    state_t0 = np.concatenate([transl_t0, R_t0_euler, hand_pose_t0_euler])  # [3+3+45=51]
    fov_x = hand_data['fov_x']

    return state_t0, hand_data[hand_side][0]['beta'], fov_x, None

def euler_traj_to_rotmat_traj(euler_traj, T):
    """
    Convert Euler angle trajectory to rotation matrix trajectory.
    
    Converts a sequence of hand poses represented as Euler angles into 
    rotation matrices suitable for MANO model input.
    
    Args:
        euler_traj (np.ndarray): Hand pose trajectory as Euler angles.
                                 Shape: [T, 45] where T is number of timesteps
                                 and 45 = 15 joints * 3 Euler angles per joint
        T (int): Number of timesteps in the trajectory
        
    Returns:
        np.ndarray: Rotation matrix trajectory. Shape: [T, 15, 3, 3]
                    where each [3, 3] block is a rotation matrix for one joint
    """
    hand_pose = euler_traj.reshape(-1, 3)  # [T*15, 3]
    pose_matrices = R.from_euler('xyz', hand_pose).as_matrix()  # [T*15, 3, 3]
    pose_matrices = pose_matrices.reshape(T, 15, 3, 3)  # [T, 15, 3, 3]

    return pose_matrices


def _hand_reconstruction_worker(args_dict, task_queue, result_queue):
    """
    Persistent worker for hand reconstruction that runs in a separate process.
    Keeps model loaded and processes multiple requests until shutdown signal.
    """
    from data.tools.hand_recon_core import Config, HandReconstructor
    
    hand_reconstructor = None
    
    try:
        # Reconstruct args object
        class ArgsObj:
            pass
        args_obj = ArgsObj()
        for key, value in args_dict.items():
            setattr(args_obj, key, value)
        
        # Initialize hand reconstructor once
        print("[HandRecon Process] Initializing hand reconstructor...")
        config = Config(args_obj)
        hand_reconstructor = HandReconstructor(config=config, device='cuda')
        print("[HandRecon Process] Hand reconstructor ready")
        
        # Signal ready
        result_queue.put({'type': 'ready'})
        
        # Process tasks in loop
        while True:
            task = task_queue.get()
            
            if task['type'] == 'shutdown':
                print("[HandRecon Process] Received shutdown signal")
                break
            
            elif task['type'] == 'reconstruct':
                try:
                    image_path = task['image_path']
                    image = cv2.imread(image_path)
                    if image is None:
                        raise ValueError(f"Failed to load image from {image_path}")
                    
                    image_list = [image]
                    recon_results = hand_reconstructor.recon(image_list)
                    
                    result_queue.put({
                        'type': 'result',
                        'success': True,
                        'data': recon_results
                    })
                    
                except Exception as e:
                    import traceback
                    result_queue.put({
                        'type': 'result',
                        'success': False,
                        'error': str(e),
                        'traceback': traceback.format_exc()
                    })
        
    except Exception as e:
        import traceback
        result_queue.put({
            'type': 'error',
            'error': str(e),
            'traceback': traceback.format_exc()
        })
    
    finally:
        # Cleanup on shutdown
        if hand_reconstructor is not None:
            del hand_reconstructor
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        print("[HandRecon Process] Cleaned up and exiting")


def _vla_inference_worker(configs_dict, task_queue, result_queue):
    """
    Persistent worker for VLA model inference that runs in a separate process.
    Keeps model loaded and processes multiple requests until shutdown signal.
    """
    from vitra.models import load_model
    from vitra.utils.data_utils import load_normalizer
    from vitra.datasets.human_dataset import pad_state_human, pad_action
    from vitra.datasets.dataset_utils import ActionFeature, StateFeature
    
    model = None
    normalizer = None
    
    try:
        # Load model and normalizer once
        print("[VLA Process] Loading VLA model...")
        model = load_model(configs_dict).cuda()
        model.eval()
        normalizer = load_normalizer(configs_dict)
        print(f"[VLA Process] VLA model ready.")
        
        # Signal ready
        result_queue.put({'type': 'ready'})
        
        # Process tasks in loop
        while True:
            task = task_queue.get()
            
            if task['type'] == 'shutdown':
                print("[VLA Process] Received shutdown signal")
                break
            
            elif task['type'] == 'predict':
                try:
                    image = task['image']
                    instruction = task['instruction']
                    state = task['state']
                    state_mask = task['state_mask']
                    action_mask = task['action_mask']
                    fov = task['fov']
                    num_ddim_steps = task.get('num_ddim_steps', 10)
                    cfg_scale = task.get('cfg_scale', 5.0)
                    sample_times = task.get('sample_times', 1)
                    
                    # Normalize state
                    norm_state = normalizer.normalize_state(state.copy())
                    
                    # Pad state and action
                    unified_action_dim = ActionFeature.ALL_FEATURES[1]  # 192
                    unified_state_dim = StateFeature.ALL_FEATURES[1]    # 212
                    
                    unified_state, unified_state_mask = pad_state_human(
                        state=norm_state,
                        state_mask=state_mask,
                        action_dim=normalizer.action_mean.shape[0],
                        state_dim=normalizer.state_mean.shape[0],
                        unified_state_dim=unified_state_dim,
                    )
                    _, unified_action_mask = pad_action(
                        actions=None,
                        action_mask=action_mask.copy(),
                        action_dim=normalizer.action_mean.shape[0],
                        unified_action_dim=unified_action_dim
                    )
                    
                    # Convert to torch and move to GPU
                    fov = torch.from_numpy(fov).unsqueeze(0)
                    unified_state = unified_state.unsqueeze(0)
                    unified_state_mask = unified_state_mask.unsqueeze(0)
                    unified_action_mask = unified_action_mask.unsqueeze(0)
                    
                    # Run inference
                    norm_action = model.predict_action(
                        image=image,
                        instruction=instruction,
                        current_state=unified_state,
                        current_state_mask=unified_state_mask,
                        action_mask_torch=unified_action_mask,
                        num_ddim_steps=num_ddim_steps,
                        cfg_scale=cfg_scale,
                        fov=fov,
                        sample_times=sample_times,
                    )
                    
                    # Extract and denormalize action
                    norm_action = norm_action[:, :, :102]
                    unnorm_action = normalizer.unnormalize_action(norm_action)
                    
                    # Convert to numpy for inter-process communication
                    if isinstance(unnorm_action, torch.Tensor):
                        unnorm_action_np = unnorm_action.cpu().numpy()
                    else:
                        unnorm_action_np = np.array(unnorm_action)
                    
                    result_queue.put({
                        'type': 'result',
                        'success': True,
                        'data': unnorm_action_np
                    })
                    
                except Exception as e:
                    import traceback
                    result_queue.put({
                        'type': 'result',
                        'success': False,
                        'error': str(e),
                        'traceback': traceback.format_exc()
                    })
        
    except Exception as e:
        import traceback
        result_queue.put({
            'type': 'error',
            'error': str(e),
            'traceback': traceback.format_exc()
        })
    
    finally:
        # Cleanup on shutdown
        if model is not None:
            del model
        if normalizer is not None:
            del normalizer
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
        print("[VLA Process] Cleaned up and exiting")


class HandReconstructionService:
    """Service wrapper for persistent hand reconstruction process"""
    
    def __init__(self, args):
        self.ctx = mp.get_context('spawn')
        self.task_queue = self.ctx.Queue()
        self.result_queue = self.ctx.Queue()
        
        # Convert args to dict for pickling
        args_dict = {
            'hawor_model_path': args.hawor_model_path,
            'detector_path': args.detector_path,
            'moge_model_name': args.moge_model_name,
            'mano_path': args.mano_path,
        }
        
        # Start persistent process
        self.process = self.ctx.Process(
            target=_hand_reconstruction_worker,
            args=(args_dict, self.task_queue, self.result_queue)
        )
        self.process.start()
        
        # Wait for ready signal
        ready_msg = self.result_queue.get()
        if ready_msg['type'] == 'ready':
            print("Hand reconstruction service initialized")
        elif ready_msg['type'] == 'error':
            raise RuntimeError(f"Failed to initialize hand reconstruction: {ready_msg['error']}")
    
    def reconstruct(self, image_path):
        """Request hand reconstruction for an image"""
        self.task_queue.put({
            'type': 'reconstruct',
            'image_path': image_path
        })
        
        result = self.result_queue.get()
        if result['type'] == 'result' and result['success']:
            return result['data']
        else:
            raise RuntimeError(f"Hand reconstruction failed: {result.get('error', 'Unknown error')}")
    
    def shutdown(self):
        """Shutdown the persistent process"""
        self.task_queue.put({'type': 'shutdown'})
        self.process.join(timeout=10)
        if self.process.is_alive():
            self.process.terminate()
            self.process.join()


class VLAInferenceService:
    """Service wrapper for persistent VLA inference process"""
    
    def __init__(self, configs):
        self.ctx = mp.get_context('spawn')
        self.task_queue = self.ctx.Queue()
        self.result_queue = self.ctx.Queue()
        
        # Start persistent process
        self.process = self.ctx.Process(
            target=_vla_inference_worker,
            args=(configs, self.task_queue, self.result_queue)
        )
        self.process.start()
        
        # Wait for ready signal
        ready_msg = self.result_queue.get()
        if ready_msg['type'] == 'ready':
            print("VLA inference service initialized")
        elif ready_msg['type'] == 'error':
            raise RuntimeError(f"Failed to initialize VLA model: {ready_msg['error']}")
    
    def predict(self, image, instruction, state, state_mask, action_mask, 
                fov, num_ddim_steps=10, cfg_scale=5.0, sample_times=1):
        """Request action prediction with state normalization and padding"""

        self.task_queue.put({
            'type': 'predict',
            'image': image,
            'instruction': instruction,
            'state': state,
            'state_mask': state_mask,
            'action_mask': action_mask,
            'fov': fov,
            'num_ddim_steps': num_ddim_steps,
            'cfg_scale': cfg_scale,
            'sample_times': sample_times,
        })
        
        result = self.result_queue.get()
        if result['type'] == 'result' and result['success']:
            # Return unnormalized action as numpy array
            return result['data']
        else:
            raise RuntimeError(f"VLA inference failed: {result.get('error', 'Unknown error')}")
    
    def shutdown(self):
        """Shutdown the persistent process"""
        self.task_queue.put({'type': 'shutdown'})
        self.process.join(timeout=10)
        if self.process.is_alive():
            self.process.terminate()
            self.process.join()


if __name__ == "__main__":
    # Set multiprocessing start method to 'spawn' to avoid CUDA fork issues
    mp.set_start_method('spawn', force=True)
    main()