Enable inference serving && fix train stability (#59)
* enable serving * lint * update * update
This commit is contained in:
Lufang Chen
@@ -0,0 +1,3 @@
|
||||
from .wall_x_policy import WallXPolicy
|
||||
|
||||
__all__ = ["WallXPolicy"]
|
||||
@@ -0,0 +1,246 @@
|
||||
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,
|
||||
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)
|
||||
|
||||
# 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"] = moe_token_types
|
||||
|
||||
# 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
|
||||
|
||||
inputs["proprioception"] = state
|
||||
inputs["agent_pos_mask"] = agent_pos_mask
|
||||
|
||||
# Add dataset name (required by model)
|
||||
inputs["dataset_names"] = obs["dataset_names"]
|
||||
|
||||
# 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:
|
||||
current_width, current_height = img_pil.size
|
||||
|
||||
# Apply smart scaling (Qwen logic)
|
||||
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 = "fast",
|
||||
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",
|
||||
}
|
||||
pred_horizon = 32
|
||||
|
||||
# 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"
|
||||
if predict_mode == "diffusion":
|
||||
assistant_output += f"{action_symbol * pred_horizon}"
|
||||
complete_text = prologue + user_message + assistant_output
|
||||
|
||||
return complete_text
|
||||
@@ -0,0 +1,173 @@
|
||||
import logging
|
||||
from typing import Dict, Any, List
|
||||
import torch
|
||||
import numpy as np
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from wall_x.serving.websocket_policy_server import BasePolicy
|
||||
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction
|
||||
from wall_x.serving.policy.utils import prepare_batch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WallXPolicy(BasePolicy):
|
||||
"""Policy wrapper for Wall-X model that implements the BasePolicy interface."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
train_config: dict,
|
||||
action_tokenizer_path: str,
|
||||
action_dim: int,
|
||||
agent_pos_dim: int,
|
||||
pred_horizon: int,
|
||||
camera_key: List[str],
|
||||
device: str = "cuda",
|
||||
dtype: str = "bfloat16",
|
||||
predict_mode: str = "fast",
|
||||
default_prompt: str | None = None,
|
||||
min_pixels: int = 4 * 28 * 28,
|
||||
max_pixels: int = 16384 * 28 * 28,
|
||||
image_factor: int = 28,
|
||||
max_length: int = 768,
|
||||
):
|
||||
"""Initialize the Wall-X policy.
|
||||
|
||||
Args:
|
||||
model_path: Path to the pretrained model checkpoint
|
||||
action_tokenizer_path: Path to the action tokenizer
|
||||
action_dim: Dimension of action space
|
||||
pred_horizon: Prediction horizon for actions
|
||||
device: Device to run model on ('cuda' or 'cpu')
|
||||
dtype: Data type for model ('bfloat16', 'float16', or 'float32')
|
||||
predict_mode: Prediction mode ('fast' or 'slow')
|
||||
default_prompt: Default text prompt for the model
|
||||
min_pixels: Minimum pixels for image resizing
|
||||
max_pixels: Maximum pixels for image resizing
|
||||
image_factor: Factor for smart resize
|
||||
max_length: Maximum sequence length for text
|
||||
"""
|
||||
logger.info(f"Loading Wall-X model from {model_path}")
|
||||
|
||||
self.model = Qwen2_5_VLMoEForAction.from_pretrained(
|
||||
model_path,
|
||||
train_config=train_config,
|
||||
action_tokenizer_path=action_tokenizer_path,
|
||||
)
|
||||
self.model.eval()
|
||||
self.model = self.model.to(device)
|
||||
|
||||
self.model = self.model.bfloat16()
|
||||
|
||||
# hard code the action dim to 20 for align to wall-x configuration
|
||||
self.fixed_action_dim = 20
|
||||
|
||||
self.action_dim = action_dim
|
||||
self.agent_pos_dim = agent_pos_dim
|
||||
self.pred_horizon = pred_horizon
|
||||
self.device = device
|
||||
self.predict_mode = predict_mode
|
||||
self.default_prompt = default_prompt
|
||||
self.camera_key = camera_key
|
||||
|
||||
# Image preprocessing config
|
||||
self.min_pixels = min_pixels
|
||||
self.max_pixels = max_pixels
|
||||
self.image_factor = image_factor
|
||||
self.max_length = max_length
|
||||
|
||||
# Load processor
|
||||
logger.info("Loading processor and tokenizer...")
|
||||
self.processor = AutoProcessor.from_pretrained(model_path, use_fast=True)
|
||||
self.processor.tokenizer.padding_side = "left"
|
||||
|
||||
# Action buffer for multi-step predictions
|
||||
self.action_buffer = []
|
||||
self.buffer_index = 0
|
||||
|
||||
logger.info(
|
||||
f"Model loaded successfully. Device: {device}, Action dim: {action_dim}, Horizon: {pred_horizon}"
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata(self) -> Dict[str, Any]:
|
||||
"""Return metadata about the policy."""
|
||||
return {
|
||||
"action_dim": self.action_dim,
|
||||
"pred_horizon": self.pred_horizon,
|
||||
"device": self.device,
|
||||
"predict_mode": self.predict_mode,
|
||||
}
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the policy state."""
|
||||
self.action_buffer = []
|
||||
self.buffer_index = 0
|
||||
logger.debug("Policy reset")
|
||||
|
||||
def infer(self, obs: Dict) -> Dict:
|
||||
"""Infer action from observation.
|
||||
|
||||
Args:
|
||||
obs: Dictionary containing:
|
||||
- 'image': Image observation (numpy array or PIL Image)
|
||||
- 'prompt': Optional text prompt
|
||||
- 'state': Optional robot state
|
||||
- Other modality-specific observations
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- 'action': Predicted action (numpy array)
|
||||
- Additional metadata
|
||||
"""
|
||||
try:
|
||||
# Need to predict new actions
|
||||
input_batch = prepare_batch(
|
||||
obs,
|
||||
self.processor,
|
||||
self.camera_key,
|
||||
self.agent_pos_dim,
|
||||
self.action_dim,
|
||||
self.pred_horizon,
|
||||
self.fixed_action_dim,
|
||||
self.max_length,
|
||||
self.image_factor,
|
||||
self.min_pixels,
|
||||
self.max_pixels,
|
||||
self.predict_mode,
|
||||
self.device,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.model(
|
||||
**input_batch,
|
||||
action_dim=(
|
||||
self.action_dim
|
||||
if self.predict_mode == "fast"
|
||||
else self.fixed_action_dim
|
||||
),
|
||||
pred_horizon=self.pred_horizon,
|
||||
mode="predict",
|
||||
predict_mode=self.predict_mode,
|
||||
)
|
||||
|
||||
if outputs["predict_action"] is None:
|
||||
predicted_actions = np.zeros(
|
||||
[1, self.pred_horizon, self.action_dim]
|
||||
).astype(np.float32)
|
||||
|
||||
predicted_actions = (
|
||||
outputs["predict_action"][:, :, : self.action_dim]
|
||||
.detach()
|
||||
.cpu()
|
||||
.to(torch.float32)
|
||||
.numpy()
|
||||
)
|
||||
|
||||
print(predicted_actions.shape)
|
||||
return {"action": predicted_actions}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during inference: {e}")
|
||||
raise
|
||||
Reference in New Issue
Block a user