[None][refactor] Cosmos3: resolve checkpoint and request settings through one parameter resolver

tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py

@@ -13,6 +13,8 @@
 import torch
 from diffusers.utils.torch_utils import randn_tensor
 
+from .defaults import EMBODIMENT_TO_DOMAIN_ID, find_closest_target_size
+
 ACTION_MODE_POLICY = "policy"
 ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics"
 ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics"
@@ -22,134 +24,6 @@
     ACTION_MODE_INVERSE_DYNAMICS,
 }
 
-EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = {
-    "no_action": 0,
-    "av": 1,
-    "camera_pose": 2,
-    "hand_pose": 3,
-    "pusht": 4,
-    "libero": 5,
-    "umi": 6,
-    "bridge_orig_lerobot": 7,
-    "droid_lerobot": 8,
-    "robomind-franka": 8,
-    "galbot": 9,
-    "robomind-franka-dual": 12,
-    "robomind-ur": 13,
-    "agibotworld": 15,
-    "agibot_gear_gripper": 15,
-    "agibot_gear_gripper_ext": 15,
-    "fractal": 20,
-}
-
-# Canonical unpadded action width per embodiment.  Widths compose the Cosmos3
-# unified action representation from shared geometric blocks: a 9-D pose (3-D
-# translation + 6-D rotation), a 1-D grasp state, and a 15-D fingertip state.
-# One arm is 9 + 1 = 10; a dual-arm setup is 20; the AgiBot humanoid is
-# 9 + 2 x (9 + 1) = 29; two-hand egocentric motion is 9 + 2 x (9 + 15) = 57.
-#
-# This is a property of the embodiment, not a tunable, so it is keyed by the
-# real domain name rather than by the sampling presets in ``defaults.py`` (where
-# several of these names share one preset).  ``libero`` is absent on purpose:
-# its width depends on the dataset's rotation space (7/10/13), so callers must
-# pass ``raw_action_dim`` explicitly.
-EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = {
-    "av": 9,
-    "camera_pose": 9,
-    "hand_pose": 57,
-    "pusht": 2,
-    "umi": 10,
-    "bridge_orig_lerobot": 10,
-    "droid_lerobot": 10,
-    "robomind-franka": 10,
-    "robomind-franka-dual": 20,
-    "robomind-ur": 10,
-    "galbot": 30,
-    "agibotworld": 29,
-    "agibot_gear_gripper": 29,
-    "agibot_gear_gripper_ext": 29,
-    "fractal": 10,
-}
-
-
-def resolve_raw_action_dim(
-    domain_name: Any = None,
-    domain_id: Any = None,
-) -> int | None:
-    """Look up the canonical action width, or None when it cannot be determined.
-
-    Resolves by name first.  A bare ``domain_id`` is only usable when every
-    embodiment sharing that id agrees on a width (true for all current ids).
-    """
-    if domain_name is not None and str(domain_name).strip():
-        return EMBODIMENT_TO_RAW_ACTION_DIM.get(str(domain_name).strip().lower())
-
-    if domain_id is None:
-        return None
-
-    widths = {
-        EMBODIMENT_TO_RAW_ACTION_DIM[name]
-        for name, mapped_id in EMBODIMENT_TO_DOMAIN_ID.items()
-        if mapped_id == int(domain_id) and name in EMBODIMENT_TO_RAW_ACTION_DIM
-    }
-    return widths.pop() if len(widths) == 1 else None
-
-
-# Canonical ``W,H`` labels; every action canvas is one of these bucket shapes.
-ACTION_ASPECT_RATIO_LABELS = ("1,1", "4,3", "3,4", "16,9", "9,16")
-
-
-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),
-    },
-}
-
-
-COSMOS3_ACTION_RESOLUTIONS = tuple(int(key) for key in sorted(VIDEO_RES_SIZE_INFO, key=int))
-
-
-def normalize_action_resolution(resolution: Any) -> int:
-    if resolution is None:
-        raise ValueError("Cosmos3 action_resolution is required for action generation.")
-    try:
-        bucket = int(resolution)
-    except (TypeError, ValueError) as exc:
-        raise ValueError(
-            f"Cosmos3 action_resolution must be an int bucket, got {resolution!r}."
-        ) from exc
-    if bucket not in COSMOS3_ACTION_RESOLUTIONS:
-        raise ValueError(
-            f"Unknown Cosmos3 action_resolution={bucket}; "
-            f"expected one of {COSMOS3_ACTION_RESOLUTIONS}."
-        )
-    return bucket
-
 
 def normalize_action_mode(mode: Any) -> str | None:
     if mode is None:
@@ -321,25 +195,6 @@ def load_action_state(action: Any = None) -> torch.Tensor:
     return tensor
 
 
-def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]:
-    key = str(resolution)
-    if key not in VIDEO_RES_SIZE_INFO:
-        raise ValueError(
-            f"Unknown Cosmos3 action resolution={resolution!r}; "
-            f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}."
-        )
-    input_ratio = h / w
-    best_size = None
-    best_diff = float("inf")
-    for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values():
-        diff = abs(input_ratio - cand_h / cand_w)
-        if diff < best_diff:
-            best_diff = diff
-            best_size = (cand_w, cand_h)
-    assert best_size is not None
-    return best_size
-
-
 ACTION_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".bmp"})
 ACTION_VIDEO_EXTENSIONS = frozenset({".mp4", ".avi"})