Enable inference serving && fix train stability (#59)
* enable serving * lint * update * update
This commit is contained in:
Lufang Chen
@@ -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
|
||||
Reference in New Issue
Block a user