fix(multimodal): stop the media loaders from importing torch

components/src/dynamo/common/multimodal/__init__.py

@@ -1,44 +1,116 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Multimodal utilities for Dynamo components."""
+"""Multimodal utilities for Dynamo components.
 
-from collections.abc import Callable
+`AsyncEncoderCache` and the embedding-transfer members import `torch`, while the
+image, audio and video loaders do not. Re-exporting all of them eagerly meant
+`from dynamo.common.multimodal import ImageLoader` pulled in the whole ML stack,
+so the loaders could not be imported or unit-tested without it. The heavy names
+resolve on first access instead (PEP 562); the loaders stay eager because they
+cost nothing.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
-from dynamo.common.constants import EmbeddingTransferMode
-from dynamo.common.multimodal.async_encoder_cache import AsyncEncoderCache
 from dynamo.common.multimodal.audio_loader import AudioLoader
-from dynamo.common.multimodal.embedding_transfer import (
-    AbstractEmbeddingReceiver,
-    AbstractEmbeddingSender,
-    LocalEmbeddingReceiver,
-    LocalEmbeddingSender,
-    NixlReadEmbeddingReceiver,
-    NixlReadEmbeddingSender,
-    NixlWriteEmbeddingReceiver,
-    NixlWriteEmbeddingSender,
-    TransferRequest,
-)
 from dynamo.common.multimodal.image_loader import ImageLoader
 from dynamo.common.multimodal.video_loader import VideoLoader
 
-EMBEDDING_SENDER_FACTORIES: dict[
-    EmbeddingTransferMode, Callable[[], AbstractEmbeddingSender]
-] = {
-    EmbeddingTransferMode.LOCAL: LocalEmbeddingSender,
-    EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingSender,
-    EmbeddingTransferMode.NIXL_READ: NixlReadEmbeddingSender,
-}
-
-EMBEDDING_RECEIVER_FACTORIES: dict[
-    EmbeddingTransferMode, Callable[[], AbstractEmbeddingReceiver]
-] = {
-    EmbeddingTransferMode.LOCAL: LocalEmbeddingReceiver,
-    EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingReceiver,
-    # [gluo FIXME] can't use pre-registered tensor as NIXL requires descriptors
-    # to be at matching size, need to overwrite nixl connect library
-    EmbeddingTransferMode.NIXL_READ: lambda: NixlReadEmbeddingReceiver(max_items=0),
-}
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from dynamo.common.constants import EmbeddingTransferMode
+    from dynamo.common.multimodal.async_encoder_cache import AsyncEncoderCache
+    from dynamo.common.multimodal.embedding_transfer import (
+        AbstractEmbeddingReceiver,
+        AbstractEmbeddingSender,
+        LocalEmbeddingReceiver,
+        LocalEmbeddingSender,
+        NixlReadEmbeddingReceiver,
+        NixlReadEmbeddingSender,
+        NixlWriteEmbeddingReceiver,
+        NixlWriteEmbeddingSender,
+        TransferRequest,
+    )
+
+_EMBEDDING_TRANSFER_EXPORTS = frozenset(
+    {
+        "AbstractEmbeddingReceiver",
+        "AbstractEmbeddingSender",
+        "LocalEmbeddingReceiver",
+        "LocalEmbeddingSender",
+        "NixlReadEmbeddingReceiver",
+        "NixlReadEmbeddingSender",
+        "NixlWriteEmbeddingReceiver",
+        "NixlWriteEmbeddingSender",
+        "TransferRequest",
+    }
+)
+
+
+def _build_sender_factories() -> (
+    dict[EmbeddingTransferMode, Callable[[], AbstractEmbeddingSender]]
+):
+    from dynamo.common.constants import EmbeddingTransferMode
+    from dynamo.common.multimodal.embedding_transfer import (
+        LocalEmbeddingSender,
+        NixlReadEmbeddingSender,
+        NixlWriteEmbeddingSender,
+    )
+
+    return {
+        EmbeddingTransferMode.LOCAL: LocalEmbeddingSender,
+        EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingSender,
+        EmbeddingTransferMode.NIXL_READ: NixlReadEmbeddingSender,
+    }
+
+
+def _build_receiver_factories() -> (
+    dict[EmbeddingTransferMode, Callable[[], AbstractEmbeddingReceiver]]
+):
+    from dynamo.common.constants import EmbeddingTransferMode
+    from dynamo.common.multimodal.embedding_transfer import (
+        LocalEmbeddingReceiver,
+        NixlReadEmbeddingReceiver,
+        NixlWriteEmbeddingReceiver,
+    )
+
+    return {
+        EmbeddingTransferMode.LOCAL: LocalEmbeddingReceiver,
+        EmbeddingTransferMode.NIXL_WRITE: NixlWriteEmbeddingReceiver,
+        # [gluo FIXME] can't use pre-registered tensor as NIXL requires descriptors
+        # to be at matching size, need to overwrite nixl connect library
+        EmbeddingTransferMode.NIXL_READ: lambda: NixlReadEmbeddingReceiver(max_items=0),
+    }
+
+
+def __getattr__(name: str) -> Any:
+    if name == "AsyncEncoderCache":
+        from dynamo.common.multimodal.async_encoder_cache import AsyncEncoderCache
+
+        value: Any = AsyncEncoderCache
+    elif name in _EMBEDDING_TRANSFER_EXPORTS:
+        from dynamo.common.multimodal import embedding_transfer
+
+        value = getattr(embedding_transfer, name)
+    elif name == "EMBEDDING_SENDER_FACTORIES":
+        value = _build_sender_factories()
+    elif name == "EMBEDDING_RECEIVER_FACTORIES":
+        value = _build_receiver_factories()
+    else:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    # Cache on the module so repeat access skips this function entirely.
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)
+
 
 __all__ = [
     "AsyncEncoderCache",

components/src/dynamo/common/tests/multimodal/test_lazy_imports.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""The package must not drag `torch` in behind the lightweight media loaders."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import pytest
+
+pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0]
+
+# A fresh interpreter is the only honest way to ask this: once any earlier test
+# has imported torch, an eager import here would go unnoticed.
+_PROBE = """
+import sys
+from dynamo.common.multimodal import {names}
+print("torch" in sys.modules)
+"""
+
+
+def _torch_imported_by(names: str) -> bool:
+    result = subprocess.run(
+        [sys.executable, "-c", _PROBE.format(names=names)],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return result.stdout.strip() == "True"
+
+
+@pytest.mark.parametrize("name", ["ImageLoader", "AudioLoader", "VideoLoader"])
+def test_media_loader_import_does_not_pull_torch(name: str) -> None:
+    assert not _torch_imported_by(
+        name
+    ), f"importing {name} pulled in torch; the loaders must stay usable without the ML stack"
+
+
+def test_heavy_members_still_import() -> None:
+    assert _torch_imported_by(
+        "AsyncEncoderCache"
+    ), "AsyncEncoderCache should still resolve, and it legitimately needs torch"