Skip to content

vllm_omni.diffusion.models.cosmos3.utils

COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION module-attribute 

COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1)

COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP module-attribute 

COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = 'first'

COSMOS3_VAE_TEMPORAL_COMPRESSION module-attribute 

COSMOS3_VAE_TEMPORAL_COMPRESSION = 4

ROBOLAB_CONCAT_VIEW_DESCRIPTION module-attribute 

ROBOLAB_CONCAT_VIEW_DESCRIPTION = "The top row is from the wrist-mounted camera. The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite sides, with the robot visible."

ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE module-attribute 

ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE = 32

ROBOLAB_DEFAULT_ACTION_SPACE module-attribute 

ROBOLAB_DEFAULT_ACTION_SPACE = 'joint_pos'

ROBOLAB_DEFAULT_CONDITIONING_FPS module-attribute 

ROBOLAB_DEFAULT_CONDITIONING_FPS = 15.0

ROBOLAB_DEFAULT_DOMAIN_NAME module-attribute 

ROBOLAB_DEFAULT_DOMAIN_NAME = 'droid_lerobot'

ROBOLAB_DEFAULT_FLOW_SHIFT module-attribute 

ROBOLAB_DEFAULT_FLOW_SHIFT = 5.0

ROBOLAB_DEFAULT_GUIDANCE_SCALE module-attribute 

ROBOLAB_DEFAULT_GUIDANCE_SCALE = 3.0

ROBOLAB_DEFAULT_IMAGE_HEIGHT module-attribute 

ROBOLAB_DEFAULT_IMAGE_HEIGHT = 540

ROBOLAB_DEFAULT_IMAGE_WIDTH module-attribute 

ROBOLAB_DEFAULT_IMAGE_WIDTH = 640

ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS module-attribute 

ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS = 4

ROBOLAB_DEFAULT_RAW_ACTION_DIM module-attribute 

ROBOLAB_DEFAULT_RAW_ACTION_DIM = 8

ROBOLAB_DEFAULT_RESOLUTION module-attribute 

ROBOLAB_DEFAULT_RESOLUTION = '480'

ROBOLAB_DEFAULT_SEED module-attribute 

ROBOLAB_DEFAULT_SEED = 0

ROBOLAB_MIDTRAIN_POSE_ACTION_DIM module-attribute 

ROBOLAB_MIDTRAIN_POSE_ACTION_DIM = 9

ROBOLAB_MIDTRAIN_RAW_ACTION_DIM module-attribute 

ROBOLAB_MIDTRAIN_RAW_ACTION_DIM = 10

VIDEO_RES_SIZE_INFO module-attribute 

VIDEO_RES_SIZE_INFO: dict[
    str, dict[str, tuple[int, int]]
] = {
    "256": {
        "1,1": (256, 256),
        "4,3": (320, 256),
        "3,4": (256, 320),
        "16,9": (320, 192),
        "9,16": (192, 320),
    },
    "480": {
        "1,1": (640, 640),
        "4,3": (736, 544),
        "3,4": (544, 736),
        "16,9": (832, 480),
        "9,16": (480, 832),
    },
    "704": {
        "1,1": (960, 960),
        "4,3": (1088, 832),
        "3,4": (832, 1088),
        "16,9": (1280, 704),
        "9,16": (704, 1280),
    },
    "720": {
        "1,1": (960, 960),
        "4,3": (1104, 832),
        "3,4": (832, 1104),
        "16,9": (1280, 720),
        "9,16": (720, 1280),
    },
}

logger module-attribute 

logger = init_logger(__name__)

RoboLabActionPostprocessInputs dataclass

action_space instance-attribute 

action_space: str

eef_pos class-attribute instance-attribute 

eef_pos: ndarray | None = None

eef_quat class-attribute instance-attribute 

eef_quat: ndarray | None = None

history_length instance-attribute 

history_length: int

RoboLabPolicyInputs dataclass

action_condition_indexes instance-attribute 

action_condition_indexes: list[int]

action_space instance-attribute 

action_space: str

action_start_frame_offset instance-attribute 

action_start_frame_offset: int

action_tensor instance-attribute 

action_tensor: Tensor

domain_id instance-attribute 

domain_id: int

flow_shift instance-attribute 

flow_shift: float

fps instance-attribute 

fps: float

guidance_scale instance-attribute 

guidance_scale: float

height instance-attribute 

height: int

history_length instance-attribute 

history_length: int

image_size instance-attribute 

image_size: Any

num_frames instance-attribute 

num_frames: int

num_inference_steps instance-attribute 

num_inference_steps: int

observation instance-attribute 

observation: dict[str, Any]

prompt instance-attribute 

prompt: str

raw_action_dim instance-attribute 

raw_action_dim: int

seed instance-attribute 

seed: int

video_tensor instance-attribute 

video_tensor: Tensor

width instance-attribute 

width: int

build_abs_pose_from_components 

build_abs_pose_from_components(*args, **kwargs) -> ndarray

build_robolab_unipc_scheduler 

build_robolab_unipc_scheduler(
    num_steps: int, shift: float, device: device
)

compose_robolab_views 

compose_robolab_views(
    obs: dict[str, Any],
) -> ndarray | None

condition_pixel_frame_count 

condition_pixel_frame_count(
    condition_frame_indexes_vision: Iterable[int],
    temporal_compression: int = COSMOS3_VAE_TEMPORAL_COMPRESSION,
) -> int

convert_midtrain_rotation 

convert_midtrain_rotation(
    value: Any, src: str, dst: str
) -> ndarray

ensure_2d_float_array 

ensure_2d_float_array(
    value: Any, key: str, width: int | None = None
) -> ndarray

ensure_gripper_array 

ensure_gripper_array(value: Any) -> ndarray

ensure_rgb_uint8_image 

ensure_rgb_uint8_image(value: Any, key: str) -> ndarray

extract_robolab_image 

extract_robolab_image(obs: dict[str, Any]) -> ndarray

extract_robolab_prompt_image 

extract_robolab_prompt_image(
    prompt_data: Any | None,
) -> ndarray | None

get_robolab_domain_id 

get_robolab_domain_id(domain_name: str) -> int

lazy_action_transform_pipeline 

lazy_action_transform_pipeline(
    max_action_dim: int,
    *,
    format_prompt_as_json: bool = False,
)

lazy_import 

lazy_import(
    module_name: str, symbol_name: str, error_message: str
)

log_robolab_action_summary 

log_robolab_action_summary(label: str, value: Any) -> None

make_robolab_action_postprocess_inputs 

make_robolab_action_postprocess_inputs(
    inputs: RoboLabPolicyInputs,
) -> RoboLabActionPostprocessInputs

next_robolab_seed 

next_robolab_seed(
    extra: dict[str, Any],
    obs: dict[str, Any],
    request_id: str | None,
) -> int

normalize_condition_frame_indexes_vision 

normalize_condition_frame_indexes_vision(
    value: Any,
) -> tuple[int, ...]

Normalize Cosmos3 vision-conditioning latent frame indexes.

normalize_condition_video_keep 

normalize_condition_video_keep(value: Any) -> str

normalize_robolab_action_space 

normalize_robolab_action_space(value: Any) -> str

pose_abs_to_rel 

pose_abs_to_rel(*args, **kwargs) -> ndarray

pose_rel_to_abs 

pose_rel_to_abs(*args, **kwargs) -> ndarray

postprocess_robolab_action 

postprocess_robolab_action(
    action: Tensor, inputs: RoboLabActionPostprocessInputs
) -> ndarray

preflight_cosmos3_action_framework_imports 

preflight_cosmos3_action_framework_imports() -> None

Validate every Cosmos Framework symbol used by action policy serving.

resize_rgb_uint8 

resize_rgb_uint8(
    image: ndarray, size: tuple[int, int]
) -> ndarray