addressed the ltx2 issue and the import issue

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>

Author: jingyu-ml

examples/diffusers/sparsity/wan22_skip_softmax.py

@@ -53,6 +53,7 @@
 import os
 
 import torch
+from datasets import load_dataset
 from diffusers import AutoencoderKLWan, WanPipeline
 from diffusers.utils import export_to_video
 
@@ -258,8 +259,6 @@ def build_sparse_config(args: argparse.Namespace, num_blocks: int) -> dict:
 
 def load_calib_prompts(calib_size: int) -> list[str]:
     """Load calibration prompts from OpenVid-1M dataset."""
-    from datasets import load_dataset
-
     dataset = load_dataset("nkp37/OpenVid-1M", split="train")
     prompts = list(dataset["caption"][:calib_size])
     print(f"Loaded {len(prompts)} calibration prompts from OpenVid-1M")

modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py

@@ -25,7 +25,6 @@
 import warnings
 
 import torch
-from ltx_core.model.transformer.attention import Attention
 
 from modelopt.torch.kernels import attention, attention_calibrate
 
@@ -171,19 +170,24 @@ def __call__(self, q, k, v, heads, mask=None):
 
 def register_ltx_triton_attention(model: torch.nn.Module) -> None:
     """Patch all ``ltx_core.Attention`` modules for Triton dispatch."""
-    warnings.warn(
-        "LTX-2 packages (ltx-core, ltx-pipelines, ltx-trainer) are provided by Lightricks "
-        "and are NOT covered by the Apache 2.0 license governing NVIDIA Model Optimizer. "
-        "You MUST comply with the LTX Community License Agreement when installing and using "
-        "LTX-2 with NVIDIA Model Optimizer. Any derivative models or fine-tuned weights from "
-        "LTX-2 (including quantized or distilled checkpoints) remain subject to the LTX "
-        "Community License Agreement, not Apache 2.0. "
-        "See: https://github.com/Lightricks/LTX-2/blob/main/LICENSE",
-        UserWarning,
-        stacklevel=2,
-    )
+    from ltx_core.model.transformer.attention import Attention
+
+    _warned = False
     for module in model.modules():
         if isinstance(module, Attention):
+            if not _warned:
+                warnings.warn(
+                    "LTX-2 packages (ltx-core, ltx-pipelines, ltx-trainer) are provided by "
+                    "Lightricks and are NOT covered by the Apache 2.0 license governing NVIDIA "
+                    "Model Optimizer. You MUST comply with the LTX Community License Agreement "
+                    "when installing and using LTX-2 with NVIDIA Model Optimizer. Any derivative "
+                    "models or fine-tuned weights from LTX-2 (including quantized or distilled "
+                    "checkpoints) remain subject to the LTX Community License Agreement, not "
+                    "Apache 2.0. See: https://github.com/Lightricks/LTX-2/blob/main/LICENSE",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                _warned = True
             fn = module.attention_function
             if not isinstance(fn, _TritonLTXAttentionWrapper):
                 module.attention_function = _TritonLTXAttentionWrapper(fn)