refactor: eager llm_worker import, add visual_gen availability check (#6165)

Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
Co-authored-by: Yuewei Na <nv-yna@users.noreply.github.com>

components/src/dynamo/trtllm/workers/__init__.py

@@ -8,6 +8,12 @@ This package contains worker initialization functions for different modalities:
 - video_diffusion_worker: Video generation using diffusion models
 
 The init_worker() function dispatches to the appropriate worker based on modality.
+
+Note on import strategy:
+- llm_worker is imported eagerly (standard dependency, always available)
+- video_diffusion_worker is imported lazily because it depends on visual_gen,
+  an optional package only available on TensorRT-LLM's feat/visual_gen branch.
+  Eager import would break text/multimodal users who don't have it installed.
 """
 
 import asyncio
@@ -16,6 +22,7 @@ import logging
 from dynamo.runtime import DistributedRuntime
 from dynamo.trtllm.constants import Modality
 from dynamo.trtllm.utils.trtllm_utils import Config
+from dynamo.trtllm.workers.llm_worker import init_llm_worker
 
 
 async def init_worker(
@@ -46,9 +53,7 @@ async def init_worker(
         # TODO: Add IMAGE_DIFFUSION support in follow-up PR
         raise ValueError(f"Unsupported diffusion modality: {modality}")
 
-    # LLM modalities (text, multimodal) use the LLM worker
-    from dynamo.trtllm.workers.llm_worker import init_llm_worker
-
+    # LLM modalities (text, multimodal)
     await init_llm_worker(runtime, config, shutdown_event)
 
 

components/src/dynamo/trtllm/workers/video_diffusion_worker.py

@@ -28,7 +28,22 @@ async def init_video_diffusion_worker(
         config: Configuration parsed from command line.
         shutdown_event: Event to signal shutdown.
     """
-    # Import diffusion-specific modules (lazy import to avoid loading heavy deps early)
+    # Check visual_gen availability early with a clear error message.
+    # visual_gen is part of TensorRT-LLM but only available on the feat/visual_gen
+    # branch — not yet in any release. Without this check, users would get a cryptic
+    # ImportError deep inside DiffusionEngine.initialize().
+    try:
+        import visual_gen  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "Video diffusion requires the 'visual_gen' package from TensorRT-LLM's "
+            "feat/visual_gen branch. Install with:\n"
+            "  git clone https://github.com/NVIDIA/TensorRT-LLM.git\n"
+            "  cd TensorRT-LLM && git checkout feat/visual_gen\n"
+            "  cd tensorrt_llm/visual_gen && pip install -e .\n"
+            "See: https://github.com/NVIDIA/TensorRT-LLM/tree/feat/visual_gen/tensorrt_llm/visual_gen"
+        ) from None
+
     from dynamo.trtllm.configs.diffusion_config import DiffusionConfig
     from dynamo.trtllm.engines.diffusion_engine import DiffusionEngine
     from dynamo.trtllm.request_handlers.video_diffusion import VideoGenerationHandler