[https://nvbugs/6437341][fix] Port sibling fix 38db241414 verbatim (shutil.which short-circuit +…

tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py

@@ -30,7 +30,7 @@
     "width": 1280,
     "num_inference_steps": 35,
     "guidance_scale": 6.0,
-    "max_sequence_length": 4096,
+    "max_sequence_length": 1024,
     "num_frames": 189,
     "frame_rate": 24.0,
 }

tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py

@@ -39,7 +39,14 @@
 from .sound_tokenizer import LatentAutoEncoderV2
 from .transformer_cosmos3 import Cosmos3VFMTransformer
 
-COSMOS3_DEFAULT_NEGATIVE_PROMPT = ""
+COSMOS3_DEFAULT_NEGATIVE_PROMPT = (
+    "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
+    "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
+    "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, "
+    "low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, "
+    "unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of "
+    "poor quality."
+)
 COSMOS3_DEFAULT_SYSTEM_PROMPT = (
     "You are a helpful assistant who will generate videos from a given prompt."
 )

tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py

@@ -351,7 +351,6 @@ def forward(
         freqs_cos: torch.Tensor,
         freqs_sin: torch.Tensor,
         timestep=None,
-        real_text_lens: Optional[list[int]] = None,
     ) -> torch.Tensor:
         """
         Args:
@@ -375,33 +374,16 @@ def forward(
         q, k = self.apply_qk_norm(q, k)
         q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
 
-        if real_text_lens is not None and batch_size > 1:
-            outs = []
-            for b in range(batch_size):
-                Lb = int(real_text_lens[b])
-                k_all_b = torch.cat([k_und[b : b + 1, :Lb], k[b : b + 1]], dim=1)
-                v_all_b = torch.cat([v_und[b : b + 1, :Lb], v[b : b + 1]], dim=1)
-                outs.append(
-                    self._attn_impl(
-                        q[b : b + 1],
-                        k_all_b,
-                        v_all_b,
-                        attention_mask=PredefinedAttentionMask.FULL,
-                        timestep=timestep,
-                    )
-                )
-            out = torch.cat(outs, dim=0)
-        else:
-            k_all = torch.cat([k_und, k], dim=1).contiguous()
-            v_all = torch.cat([v_und, v], dim=1).contiguous()
-
-            out = self._attn_impl(
-                q,
-                k_all,
-                v_all,
-                attention_mask=PredefinedAttentionMask.FULL,
-                timestep=timestep,
-            )
+        k_all = torch.cat([k_und, k], dim=1).contiguous()
+        v_all = torch.cat([v_und, v], dim=1).contiguous()
+
+        out = self._attn_impl(
+            q,
+            k_all,
+            v_all,
+            attention_mask=PredefinedAttentionMask.FULL,
+            timestep=timestep,
+        )
 
         return self.to_out[0](out)
 
@@ -521,7 +503,6 @@ def forward(
         v_und: torch.Tensor,
         freqs: Tuple[torch.Tensor, torch.Tensor],
         timestep=None,
-        real_text_lens: Optional[list[int]] = None,
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
@@ -534,7 +515,6 @@ def forward(
             freqs_cos=cos,
             freqs_sin=sin,
             timestep=timestep,
-            real_text_lens=real_text_lens,
         )
         hidden_states = residual + hidden_states
 
@@ -1016,7 +996,6 @@ def forward(
         T, H, W = video_shape
         Hp, Wp, _, _ = self._pad_to_patch_size(H, W)
         max_real_len = text_mask.sum(dim=1).max().item()
-        real_text_lens = text_mask.sum(dim=1).tolist()
 
         hidden_gen = self.vae2llm(self.patchify(hidden_states, T, H, W))
 
@@ -1113,22 +1092,13 @@ def forward(
             if not self.sharder.is_active:
                 k_und = k_und[:, :max_real_len]
                 v_und = v_und[:, :max_real_len]
-                hidden_gen = layer(
-                    hidden_gen,
-                    k_und,
-                    v_und,
-                    freqs_gen,
-                    timestep=timestep,
-                    real_text_lens=real_text_lens,
-                )
-            else:
-                hidden_gen = layer(
-                    hidden_gen,
-                    k_und,
-                    v_und,
-                    freqs_gen,
-                    timestep=timestep,
-                )
+            hidden_gen = layer(
+                hidden_gen,
+                k_und,
+                v_und,
+                freqs_gen,
+                timestep=timestep,
+            )
 
         hidden_gen = self.sharder.gather(hidden_gen, dim=1, unpad_to=S_gen)
 

tests/integration/defs/examples/visual_gen/test_visual_gen.py

@@ -235,9 +235,38 @@ def _visual_gen_deps(llm_venv):
     """Install av + diffusers + ffmpeg once per session (shared by all video-gen fixtures)."""
     llm_venv.run_cmd(["-m", "pip", "install", "av"])
     llm_venv.run_cmd(["-m", "pip", "install", "diffusers>=0.37.0"])
-    # Install ffmpeg system package required by save_video() for MP4 encoding
-    check_call(["apt-get", "update", "-y"], shell=False)
-    check_call(["apt-get", "install", "-y", "ffmpeg"], shell=False)
+    # opencv was removed from the container image by #16206; the LPIPS eval
+    # script imports cv2 lazily for the video path (see scripts/visualgen_eval/
+    # visual_gen_lpips_score_eval.py::_get_cv2), so install the headless wheel.
+    llm_venv.run_cmd(["-m", "pip", "install", "opencv-python-headless"])
+    _ensure_ffmpeg_available(llm_venv)
+
+
+def _ensure_ffmpeg_available(llm_venv):
+    """Put an ffmpeg binary on PATH for save_video()'s MP4 encoder.
+
+    Unprivileged CI containers (LOCAL_USER=1) fail apt-get with exit 100 on
+    /var/lib/apt/lists/partial; fall back to the imageio-ffmpeg wheel, which
+    bundles a static libx264-enabled ffmpeg 7.x binary needing no root.
+    """
+    if shutil.which("ffmpeg"):
+        return
+
+    try:
+        check_call(["apt-get", "update", "-y"], shell=False)
+        check_call(["apt-get", "install", "-y", "ffmpeg"], shell=False)
+    except subprocess.CalledProcessError:
+        llm_venv.run_cmd(["-m", "pip", "install", "imageio-ffmpeg"])
+        import imageio_ffmpeg
+
+        target_dir = os.path.expanduser("~/.local/bin")
+        os.makedirs(target_dir, exist_ok=True)
+        target = os.path.join(target_dir, "ffmpeg")
+        if os.path.lexists(target):
+            os.remove(target)
+        os.symlink(imageio_ffmpeg.get_ffmpeg_exe(), target)
+
+    assert shutil.which("ffmpeg"), "ffmpeg install did not put a binary on PATH"
 
 
 @pytest.fixture(scope="session")

tests/integration/test_lists/waives.txt

@@ -190,7 +190,6 @@ examples/test_ray.py::test_llm_inference_distributed_ray[tp2pp2] SKIP (https://n
 examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/5612502)
 examples/test_whisper.py::test_llm_whisper_general[large-v3-disable_gemm_plugin-disable_attention_plugin-disable_weight_only-float16-nb:1-use_python_runtime] SKIP (https://nvbugs/5244570)
 examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815)
-examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341)
 examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[attn2d_2x2] SKIP (https://nvbugs/6272644)
 examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2] SKIP (https://nvbugs/6272644)
 examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2_attn2d_2x1] SKIP (https://nvbugs/6272644)