merge main

Author: Huy Vu2

.github/actions/test-template/action.yml

@@ -140,7 +140,7 @@ runs:
     - name: Checkout repository
       uses: actions/checkout@v2
       with:
-        path: VFM
+        path: DFM
 
     - name: Start container
       shell: bash
@@ -164,7 +164,7 @@ runs:
           --env HYDRA_FULL_ERROR=1 \
           --env HF_HOME=/home/TestData/HF_HOME \
           --env RUN_ID=${{ github.run_id }} \
-          --volume $(pwd)/VFM:/workspace \
+          --volume $(pwd)/DFM:/workspace \
           --volume $MNT_PATH/TestData:/home/TestData \
           nemoci.azurecr.io/${{ inputs.image }}:${{ github.run_id }} \
           bash -c "sleep $(( ${{ inputs.timeout }} * 60 + 60))"

dfm/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from dfm.package_info import (
+    __contact_emails__,
+    __contact_names__,
+    __description__,
+    __download_url__,
+    __homepage__,
+    __keywords__,
+    __license__,
+    __package_name__,
+    __repository_url__,
+    __shortversion__,
+    __version__,
+)

nemo_vfm/package_info.py dfm/package_info.pynemo_vfm/package_info.py renamed to dfm/package_info.py

@@ -24,12 +24,12 @@
 __shortversion__ = ".".join(map(str, VERSION[:3]))
 __version__ = ".".join(map(str, VERSION[:3])) + "".join(VERSION[3:])
 
-__package_name__ = "nemo_vfm"
+__package_name__ = "dfm"
 __contact_names__ = "NVIDIA"
 __contact_emails__ = "nemo-toolkit@nvidia.com"
-__homepage__ = "https://github.com/NVIDIA-NeMo/NeMo-VFM"
-__repository_url__ = "https://github.com/NVIDIA-NeMo/NeMo-VFM"
-__download_url__ = "https://github.com/NVIDIA-NeMo/NeMo-VFM/releases"
-__description__ = "NeMo VFM"
+__homepage__ = "https://github.com/NVIDIA-NeMo/NeMo-DFM"
+__repository_url__ = "https://github.com/NVIDIA-NeMo/NeMo-DFM"
+__download_url__ = "https://github.com/NVIDIA-NeMo/NeMo-DFM/releases"
+__description__ = "NeMo DFM"
 __license__ = "Apache2"
 __keywords__ = "deep learning, machine learning, gpu, NLP, pytorch, torch"

…on/models/dit/action_control/__init__.py dfm/src/common/tokenizers/__init__ .pynemo_vfm/diffusion/models/dit/action_control/__init__.py renamed to dfm/src/common/tokenizers/__init__ .py nemo_vfm/diffusion/models/dit/action_control/__init__.py renamed to dfm/src/common/tokenizers/__init__ .py

@@ -17,49 +17,95 @@
 
 """A library for Causal Video Tokenizer inference."""
 
-from typing import Any
+from pathlib import Path
 
 import numpy as np
 import torch
-from cosmos1.models.tokenizer.inference.utils import (
-    load_decoder_model,
-    load_encoder_model,
-    load_model,
+from huggingface_hub import get_token as get_hf_token
+from huggingface_hub import hf_hub_download
+from tqdm import tqdm
+
+from dfm.src.common.tokenizers.cosmos.cosmos1.video_tokenizer_utils import (
+    load_jit_model,
     numpy2tensor,
     pad_video_batch,
     tensor2numpy,
     unpad_video_batch,
 )
-from tqdm import tqdm
 
 
 class CausalVideoTokenizer(torch.nn.Module):
     def __init__(
         self,
-        checkpoint: str = None,
-        checkpoint_enc: str = None,
-        checkpoint_dec: str = None,
-        tokenizer_config: dict[str, Any] = None,
+        checkpoint_dir: str = None,
+        load_full_model: bool = True,
+        load_dec_model: bool = True,
+        load_enc_model: bool = True,
         device: str = "cuda",
         dtype: str = "bfloat16",
     ) -> None:
         super().__init__()
+
+        checkpoint = Path(checkpoint_dir)
+        self._full_model_path = str(checkpoint / "autoencoder.jit")
+        self._enc_model_path = str(checkpoint / "encoder.jit")
+        self._dec_model_path = str(checkpoint / "decoder.jit")
+        self._dtype = dtype
         self._device = device
-        self._dtype = getattr(torch, dtype)
-        self._full_model = (
-            load_model(checkpoint, tokenizer_config, device).to(self._dtype) if checkpoint is not None else None
+
+        self._full_model = load_jit_model(self._full_model_path, self._device) if load_full_model else None
+        self._enc_model = load_jit_model(self._enc_model_path, self._device) if load_enc_model else None
+        self._dec_model = load_jit_model(self._dec_model_path, self._device) if load_dec_model else None
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        tokenizer_type="Cosmos-Tokenizer-DV4x8x8",
+        load_encoder=True,
+        load_decoder=True,
+        load_full_model=False,
+        use_pytorch=False,
+        dtype="bfloat16",
+    ):
+        cls._hf_model_name = f"nvidia/{tokenizer_type}"
+
+        # Requires setting HF_TOKEN env variable
+        hf_token = get_hf_token()
+
+        full_model_path = hf_hub_download(
+            repo_id=cls._hf_model_name,
+            filename="autoencoder.jit",
+            token=hf_token,
         )
-        self._enc_model = (
-            load_encoder_model(checkpoint_enc, tokenizer_config, device).to(self._dtype)
-            if checkpoint_enc is not None
-            else None
+
+        _ = hf_hub_download(
+            repo_id=cls._hf_model_name,
+            filename="encoder.jit",
+            token=hf_token,
         )
-        self._dec_model = (
-            load_decoder_model(checkpoint_dec, tokenizer_config, device).to(self._dtype)
-            if checkpoint_dec is not None
-            else None
+
+        _ = hf_hub_download(
+            repo_id=cls._hf_model_name,
+            filename="decoder.jit",
+            token=hf_token,
         )
 
+        # No need to load in encoder and decoder with full model loaded
+        if load_full_model:
+            load_encoder = False
+            load_decoder = False
+
+        # Assumes HF downloads all files to same local dir
+        ckpt_dir = str(Path(full_model_path).parent)
+        args = {
+            "checkpoint_dir": ckpt_dir,
+            "dtype": dtype,
+            "load_enc_model": load_encoder,
+            "load_dec_model": load_decoder,
+            "load_full_model": load_full_model,
+        }
+        return cls(**args)
+
     @torch.no_grad()
     def autoencode(self, input_tensor: torch.Tensor) -> torch.Tensor:
         """Reconstrcuts a batch of video tensors after embedding into a latent.

…vfm/diffusion/sampler/cosmos/__init__.py …src/common/tokenizers/cosmos/__init__.pynemo_vfm/diffusion/sampler/cosmos/__init__.py renamed to dfm/src/common/tokenizers/cosmos/__init__.py nemo_vfm/diffusion/sampler/cosmos/__init__.py renamed to dfm/src/common/tokenizers/cosmos/__init__.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for the inference libraries."""
+
+import numpy as np
+import torch
+
+
+_DTYPE, _DEVICE = torch.bfloat16, "cuda"
+_UINT8_MAX_F = float(torch.iinfo(torch.uint8).max)
+_SPATIAL_ALIGN = 16
+_TEMPORAL_ALIGN = 8
+
+
+def load_jit_model(jit_filepath: str = None, device: str = "cuda") -> torch.jit.ScriptModule:
+    """Loads a torch.jit.ScriptModule from a filepath.
+
+    Args:
+        jit_filepath: The filepath to the JIT-compiled model.
+        device: The device to load the model onto, default=cuda.
+    Returns:
+        The JIT compiled model loaded to device and on eval mode.
+    """
+    model = torch.jit.load(jit_filepath, map_location=device)
+    return model.eval().to(device)
+
+
+def numpy2tensor(
+    input_image: np.ndarray,
+    dtype: torch.dtype = _DTYPE,
+    device: str = _DEVICE,
+    range_min: int = -1,
+) -> torch.Tensor:
+    """Converts image(dtype=np.uint8) to `dtype` in range [0..255].
+
+    Args:
+        input_image: A batch of images in range [0..255], BxHxWx3 layout.
+    Returns:
+        A torch.Tensor of layout Bx3xHxW in range [-1..1], dtype.
+    """
+    ndim = input_image.ndim
+    indices = list(range(1, ndim))[-1:] + list(range(1, ndim))[:-1]
+    image = input_image.transpose((0,) + tuple(indices)) / _UINT8_MAX_F
+    if range_min == -1:
+        image = 2.0 * image - 1.0
+    return torch.from_numpy(image).to(dtype).to(device)
+
+
+def tensor2numpy(input_tensor: torch.Tensor, range_min: int = -1) -> np.ndarray:
+    """Converts tensor in [-1,1] to image(dtype=np.uint8) in range [0..255].
+
+    Args:
+        input_tensor: Input image tensor of Bx3xHxW layout, range [-1..1].
+    Returns:
+        A numpy image of layout BxHxWx3, range [0..255], uint8 dtype.
+    """
+    if range_min == -1:
+        input_tensor = (input_tensor.float() + 1.0) / 2.0
+    ndim = input_tensor.ndim
+    output_image = input_tensor.clamp(0, 1).cpu().numpy()
+    output_image = output_image.transpose((0,) + tuple(range(2, ndim)) + (1,))
+    return (output_image * _UINT8_MAX_F + 0.5).astype(np.uint8)
+
+
+def pad_video_batch(
+    batch: np.ndarray,
+    temporal_align: int = _TEMPORAL_ALIGN,
+    spatial_align: int = _SPATIAL_ALIGN,
+) -> tuple[np.ndarray, list[int]]:
+    """Pads a batch of videos to be divisible by `temporal_align` or `spatial_align`.
+
+    Zero pad spatially. Reflection pad temporally to handle causality better.
+    Args:
+        batch: The batch of videos to pad., layout BxFxHxWx3, in any range.
+        align: The alignment to pad to.
+    Returns:
+        The padded batch and the crop region.
+    """
+    num_frames, height, width = batch.shape[-4:-1]
+    align = spatial_align
+    height_to_pad = (align - height % align) if height % align != 0 else 0
+    width_to_pad = (align - width % align) if width % align != 0 else 0
+
+    align = temporal_align
+    frames_to_pad = (align - (num_frames - 1) % align) if (num_frames - 1) % align != 0 else 0
+
+    crop_region = [
+        frames_to_pad >> 1,
+        height_to_pad >> 1,
+        width_to_pad >> 1,
+        num_frames + (frames_to_pad >> 1),
+        height + (height_to_pad >> 1),
+        width + (width_to_pad >> 1),
+    ]
+    batch = np.pad(
+        batch,
+        (
+            (0, 0),
+            (0, 0),
+            (height_to_pad >> 1, height_to_pad - (height_to_pad >> 1)),
+            (width_to_pad >> 1, width_to_pad - (width_to_pad >> 1)),
+            (0, 0),
+        ),
+        mode="constant",
+    )
+    batch = np.pad(
+        batch,
+        (
+            (0, 0),
+            (frames_to_pad >> 1, frames_to_pad - (frames_to_pad >> 1)),
+            (0, 0),
+            (0, 0),
+            (0, 0),
+        ),
+        mode="edge",
+    )
+    return batch, crop_region
+
+
+def unpad_video_batch(batch: np.ndarray, crop_region: list[int]) -> np.ndarray:
+    """Unpads video with `crop_region`.
+
+    Args:
+        batch: A batch of numpy videos, layout BxFxHxWxC.
+        crop_region: [f1,y1,x1,f2,y2,x2] first, top, left, last, bot, right crop indices.
+
+    Returns:
+        np.ndarray: Cropped numpy video, layout BxFxHxWxC.
+    """
+    assert len(crop_region) == 6, "crop_region should be len of 6."
+    f1, y1, x1, f2, y2, x2 = crop_region
+    return batch[..., f1:f2, y1:y2, x1:x2, :]