VersatIL 0.3.0__py3-none-any.whl

versatil/__init__.py

@@ -0,0 +1,85 @@
+"""VersatIL library."""
+
+import logging
+import os
+import warnings
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Set cache dirs before importing transformers/timm — they read
+CACHE_DIR = Path(os.environ["VERSATIL_CACHE_DIR"])
+
+
+def setup_cache_directories():
+    """Configure cache directories for model downloads."""
+    os.environ["HF_HOME"] = str(CACHE_DIR / "huggingface")
+    os.environ["HF_HUB_CACHE"] = str(CACHE_DIR / "huggingface" / "hub")
+    os.environ["TORCH_HOME"] = str(CACHE_DIR / "torch")
+    for cache_path in [
+        CACHE_DIR / "huggingface" / "transformers",
+        CACHE_DIR / "huggingface" / "hub",
+        CACHE_DIR / "torch" / "hub",
+    ]:
+        cache_path.mkdir(parents=True, exist_ok=True)
+
+
+setup_cache_directories()
+
+import transformers
+
+from versatil.quantization.torch_patches import patch_pt2e_python314
+
+patch_pt2e_python314()
+
+logging.getLogger("timm").setLevel(logging.ERROR)
+logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+transformers.logging.set_verbosity_error()
+
+warnings.filterwarnings(
+    "ignore",
+    category=UserWarning,
+    module="pydantic._internal._generate_schema",
+)
+warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
+warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
+warnings.filterwarnings("ignore", category=FutureWarning, module="timm")
+warnings.filterwarnings("ignore", category=UserWarning, module="timm")
+warnings.filterwarnings("ignore", category=UserWarning, module="albumentations")
+warnings.filterwarnings("ignore", category=UserWarning, module="hydra")
+warnings.filterwarnings("ignore", message="Trying to infer the `batch_size`")
+warnings.filterwarnings("ignore", category=SyntaxWarning, module="geomloss")
+warnings.filterwarnings("ignore", category=SyntaxWarning, module="torchao")
+warnings.filterwarnings(
+    "ignore",
+    message="The given buffer is not writable",
+    category=UserWarning,
+)
+warnings.filterwarnings(
+    "ignore",
+    message=".*self.log.*self.trainer.*not registered.*",
+    module="pytorch_lightning",
+)
+warnings.filterwarnings(
+    "ignore",
+    message=".*tensorboardX.*has been removed.*",
+    module="pytorch_lightning",
+)
+warnings.filterwarnings(
+    "ignore",
+    message=".*isinstance.*LeafSpec.*is deprecated.*",
+    module="pytorch_lightning",
+)
+warnings.filterwarnings(
+    "ignore",
+    message="Checkpoint directory.*exists and is not empty",
+    module="pytorch_lightning",
+)
+warnings.filterwarnings(
+    "ignore",
+    message="crc32c usage is deprecated",
+    module="numcodecs.*",
+)

versatil/common/__init__.py

@@ -0,0 +1 @@
+"""Common functionalities and utilities shared across the codebase."""

versatil/common/argparse_compat.py

@@ -0,0 +1,50 @@
+"""Python 3.14 argparse compatibility shim for Hydra 1.3.2.
+
+Python 3.14 added ``_ActionsContainer._check_help`` which eagerly validates
+``help=`` arguments at ``add_argument`` time by calling
+``formatter._expand_help(action)``. That path performs ``'%' in help_string``,
+which assumes ``help`` is a ``str``. Hydra 1.3.2 passes a
+``LazyCompletionHelp`` instance (defined inside
+``hydra._internal.utils.get_args_parser``) that only implements ``__repr__``,
+so the eager check raises ``ValueError('badly formed help string')`` and
+prevents any Hydra endpoint from starting on Python 3.14.
+
+This module restores the pre-3.14 deferred behavior only for non-string
+``help`` values: string help strings still go through the stock check, so
+genuinely malformed strings are still caught.
+
+Import this module once before ``import hydra``.
+
+Upstream status
+---------------
+- Hydra issue: https://github.com/facebookresearch/hydra/issues/3121
+  (open, same traceback, unresolved in 1.3.2)
+- Hydra fix: https://github.com/facebookresearch/hydra/pull/3090
+  (merged 2025-10-28 into ``main``, targets ``1.4.0.dev``, not on PyPI).
+  The upstream fix uses the same approach — temporarily disabling
+  ``argparse.ArgumentParser._check_help``.
+- Release tracker: https://github.com/facebookresearch/hydra/issues/3125
+  (asking for a 1.4 release; no release yet)
+- CPython: the eager ``_check_help`` was added in CPython PR #141940 and is
+  not flagged as a breaking change in the Python 3.14 "What's New" notes.
+
+This shim can be removed once ``hydra-core >= 1.4`` is released on PyPI
+and the project bumps the pin in ``pyproject.toml``.
+"""
+
+import argparse
+
+
+def _install() -> None:
+    original_check_help = argparse._ActionsContainer._check_help
+
+    def _check_help_string_only(
+        self: argparse._ActionsContainer, action: argparse.Action
+    ) -> None:
+        if action.help is None or isinstance(action.help, str):
+            original_check_help(self, action)
+
+    argparse._ActionsContainer._check_help = _check_help_string_only  # type: ignore[method-assign]
+
+
+_install()

versatil/common/dict_of_tensor_mixin.py

@@ -0,0 +1,48 @@
+"""Mixin for nn.Module that stores and loads nested ParameterDict state."""
+
+import torch
+import torch.nn as nn
+
+
+class DictOfTensorMixin(nn.Module):
+    def __init__(self, params_dict=None):
+        super().__init__()
+        if params_dict is None:
+            params_dict = nn.ParameterDict()
+        self.params_dict = params_dict
+        self.register_buffer("_device_tracker", torch.zeros(1))
+
+    @property
+    def device(self):
+        return self._device_tracker.device
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        def dfs_add(dest, keys, value: torch.Tensor):
+            if len(keys) == 1:
+                dest[keys[0]] = value
+                return
+
+            if keys[0] not in dest:
+                dest[keys[0]] = nn.ParameterDict()
+            dfs_add(dest[keys[0]], keys[1:], value)
+
+        def load_dict(state_dict, prefix):
+            out_dict = nn.ParameterDict()
+            for key, value in state_dict.items():
+                if key.startswith(prefix):
+                    param_keys = key[len(prefix) :].split(".")[1:]
+                    dfs_add(out_dict, param_keys, value.clone())
+            return out_dict
+
+        self.params_dict = load_dict(state_dict, prefix + "params_dict")
+        self.params_dict.requires_grad_(False)
+        return

versatil/common/logging.py

@@ -0,0 +1,16 @@
+"""Logging utilities for VersatIL endpoints."""
+
+import logging
+
+LOG_FORMAT = "%(asctime)s %(module)s %(levelname)s %(message)s"
+
+
+def override_log_format() -> None:
+    """Replace the log formatter on all root handlers.
+
+    Call after Hydra's ``@hydra.main`` has configured logging,
+    so the format shows the source module instead of ``root``.
+    """
+    formatter = logging.Formatter(LOG_FORMAT)
+    for handler in logging.root.handlers:
+        handler.setFormatter(formatter)

versatil/common/module_attr_mixin.py

@@ -0,0 +1,17 @@
+"""Mixin for nn.Module that exposes device and dtype properties."""
+
+import torch.nn as nn
+
+
+class ModuleAttrMixin(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._dummy_variable = nn.Parameter()
+
+    @property
+    def device(self):
+        return next(iter(self.parameters())).device
+
+    @property
+    def dtype(self):
+        return next(iter(self.parameters())).dtype

versatil/common/omegaconf_ops.py

@@ -0,0 +1,33 @@
+"""OmegaConf utility operations."""
+
+from collections.abc import Mapping
+from typing import Any
+
+from omegaconf import OmegaConf
+
+
+def resolve_dict_keys(d: Mapping[Any, Any]) -> dict[Any, Any]:
+    """Resolve any OmegaConf interpolations in dictionary keys recursively.
+
+    OmegaConf doesn't resolve interpolations in dict keys by default.
+    This function resolves keys like '${cameras:AGENTVIEW}' to 'agentview_rgb'.
+
+    Args:
+        d: Dictionary with potentially unresolved interpolation keys.
+
+    Returns:
+        New dictionary with resolved keys.
+    """
+    resolved = {}
+    for key, value in d.items():
+        if isinstance(key, str) and key.startswith("${") and key.endswith("}"):
+            temp_cfg = OmegaConf.create({"_key": key})
+            OmegaConf.resolve(temp_cfg)
+            resolved_key = OmegaConf.select(temp_cfg, "_key")
+        else:
+            resolved_key = key
+        resolved_value = (
+            resolve_dict_keys(value) if isinstance(value, Mapping) else value
+        )
+        resolved[resolved_key] = resolved_value
+    return resolved

versatil/common/set_cache_dir.py

@@ -0,0 +1,21 @@
+"""Configuration of HuggingFace and Torch cache directories."""
+
+import os
+from pathlib import Path
+
+
+def setup_cache_directories(cache_dir: str | Path):
+    """Configure cache directories for model downloads before any run."""
+    # Convert to Path for consistent handling
+    cache_path = Path(cache_dir)
+
+    os.environ["HF_HOME"] = str(cache_path / "huggingface")
+    os.environ["HF_HUB_CACHE"] = str(cache_path / "huggingface" / "hub")
+    os.environ["TORCH_HOME"] = str(cache_path / "torch")
+
+    for dir_path in [
+        cache_path / "huggingface" / "transformers",
+        cache_path / "huggingface" / "hub",
+        cache_path / "torch" / "hub",
+    ]:
+        dir_path.mkdir(parents=True, exist_ok=True)

versatil/common/tensor_ops.py

@@ -0,0 +1,126 @@
+"""Tensor container utilities used across VersatIL."""
+
+from collections import OrderedDict
+from collections.abc import Callable
+
+import torch
+
+type NestedDictionary[Leaf] = dict[str, Leaf | NestedDictionary[Leaf]]
+type TensorTree = (
+    torch.Tensor
+    | str
+    | None
+    | dict[str, TensorTree]
+    | OrderedDict[str, TensorTree]
+    | list[TensorTree]
+    | tuple[TensorTree, ...]
+)
+type TensorTreeHandlerMap = dict[type, Callable[[TensorTree], TensorTree]]
+
+
+def dict_apply[LeafInput, LeafOutput](
+    data: NestedDictionary[LeafInput],
+    transform: Callable[[LeafInput], LeafOutput],
+) -> NestedDictionary[LeafOutput]:
+    """Apply a transform to every leaf in a nested dictionary.
+
+    Args:
+        data: Nested dictionary whose leaves should be transformed.
+        transform: Function applied to each leaf value.
+
+    Returns:
+        Nested dictionary with the same keys and transformed leaf values.
+    """
+    result = {}
+    for key, value in data.items():
+        if isinstance(value, dict):
+            result[key] = dict_apply(data=value, transform=transform)
+        else:
+            result[key] = transform(value)
+    return result
+
+
+def tensor_to_str(tensor: torch.Tensor) -> str:
+    """Format a tensor as a compact numeric list for logs.
+
+    Args:
+        tensor: Tensor to format. Scalars are treated as one-element tensors.
+
+    Returns:
+        String with each value formatted to three significant digits.
+    """
+    flat_tensor = tensor.flatten()
+    return "[" + ", ".join([f"{value.item():.3g}" for value in flat_tensor]) + "]"
+
+
+def recursive_dict_list_tuple_apply(
+    data: TensorTree,
+    type_handler_map: TensorTreeHandlerMap,
+) -> TensorTree:
+    """Apply type-specific handlers to leaves inside nested containers.
+
+    Args:
+        data: Tensor tree containing dictionaries, sequences, and supported leaves.
+        type_handler_map: Mapping from leaf type to its transform function.
+
+    Returns:
+        Tensor tree with handlers applied to matching leaves.
+
+    Raises:
+        ValueError: If a container type is passed as a leaf handler.
+        NotImplementedError: If a leaf value has no matching handler.
+    """
+    if isinstance(data, (dict, OrderedDict)):
+        if dict in type_handler_map:
+            raise ValueError("dict cannot be handled by type_func_dict")
+        return OrderedDict(
+            [
+                (
+                    key,
+                    recursive_dict_list_tuple_apply(
+                        data=value,
+                        type_handler_map=type_handler_map,
+                    ),
+                )
+                for key, value in data.items()
+            ]
+        )
+
+    if isinstance(data, (list, tuple)):
+        if list in type_handler_map or tuple in type_handler_map:
+            raise ValueError("list/tuple cannot be handled by type_func_dict")
+        return type(data)(
+            [
+                recursive_dict_list_tuple_apply(
+                    data=value,
+                    type_handler_map=type_handler_map,
+                )
+                for value in data
+            ]
+        )
+
+    for leaf_type, handler in type_handler_map.items():
+        if isinstance(data, leaf_type):
+            return handler(data)
+
+    raise NotImplementedError(f"Unsupported type: {type(data)}")
+
+
+def to_device(data: TensorTree, device: torch.device | str) -> TensorTree:
+    """Move every tensor in a nested container to a device.
+
+    Args:
+        data: Tensor tree containing tensors and pass-through string or None leaves.
+        device: Destination device.
+
+    Returns:
+        Tensor tree with all tensor leaves moved to the destination device.
+    """
+    return recursive_dict_list_tuple_apply(
+        data=data,
+        type_handler_map={
+            torch.Tensor: lambda tensor: tensor.to(device),
+            str: lambda value: value,
+            type(None): lambda value: value,
+        },
+    )