Fix Helios V2V latent prep crash and LTX long scheduler desync on skipped steps

src/diffusers/pipelines/helios/pipeline_helios.py

@@ -390,6 +390,9 @@ def prepare_video_latents(
     ) -> torch.Tensor:
         device = device or self._execution_device
         video = video.to(device=device, dtype=self.vae.dtype)
+        first_frame = video[:, :, 0:1, :, :]
+        first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator)
+        first_frame_latent = (first_frame_latent - latents_mean) * latents_std
         if latents is None:
             num_frames = video.shape[2]
             min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1
@@ -403,10 +406,6 @@ def prepare_video_latents(
             total_valid_frames = num_chunks * min_frames
             start_frame = num_frames - total_valid_frames
 
-            first_frame = video[:, :, 0:1, :, :]
-            first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator)
-            first_frame_latent = (first_frame_latent - latents_mean) * latents_std
-
             latents_chunks = []
             for i in range(num_chunks):
                 chunk_start = start_frame + i * min_frames
@@ -694,6 +693,8 @@ def __call__(
             )
 
         if video_latents is not None and add_noise_to_video_latents:
+            if image_latents is None:
+                image_latents = video_latents[:, :, 0:1, :, :]
             image_noise_sigma = (
                 torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min)
                 + image_noise_sigma_min

src/diffusers/pipelines/helios/pipeline_helios_pyramid.py

@@ -412,6 +412,9 @@ def prepare_video_latents(
     ) -> torch.Tensor:
         device = device or self._execution_device
         video = video.to(device=device, dtype=self.vae.dtype)
+        first_frame = video[:, :, 0:1, :, :]
+        first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator)
+        first_frame_latent = (first_frame_latent - latents_mean) * latents_std
         if latents is None:
             num_frames = video.shape[2]
             min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1
@@ -425,10 +428,6 @@ def prepare_video_latents(
             total_valid_frames = num_chunks * min_frames
             start_frame = num_frames - total_valid_frames
 
-            first_frame = video[:, :, 0:1, :, :]
-            first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator)
-            first_frame_latent = (first_frame_latent - latents_mean) * latents_std
-
             latents_chunks = []
             for i in range(num_chunks):
                 chunk_start = start_frame + i * min_frames
@@ -770,6 +769,8 @@ def __call__(
             )
 
         if video_latents is not None and add_noise_to_video_latents:
+            if image_latents is None:
+                image_latents = video_latents[:, :, 0:1, :, :]
             image_noise_sigma = (
                 torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min)
                 + image_noise_sigma_min

src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py

@@ -1337,6 +1337,8 @@ def __call__(
                             )
 
                     # Use global timestep for scheduling, but apply suppressive blending with hard-condition tokens (e.g., first frame) after step to avoid brightness/flicker due to time misalignment
+                    # Re-sync step index with the current loop timestep when earlier steps were skipped.
+                    self.scheduler._step_index = None
                     latents_packed = self.scheduler.step(
                         noise_pred, t, latents_packed, generator=local_gen, return_dict=False
                     )[0]

tests/pipelines/helios/test_helios.py

@@ -139,6 +139,33 @@ def test_inference(self):
         generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
         self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3))
 
+    def test_prepare_video_latents_with_precomputed_latents(self):
+        components = self.get_dummy_components()
+        pipe = HeliosPipeline(**components)
+        pipe.to(torch_device)
+        num_latent_frames_per_chunk = 9
+        min_frames = (num_latent_frames_per_chunk - 1) * pipe.vae_scale_factor_temporal + 1
+        video = torch.randn(1, 3, min_frames, 16, 16, device=torch_device)
+        latents_mean = torch.zeros(
+            1, pipe.vae.config.z_dim, 1, 1, 1, device=torch_device, dtype=pipe.vae.dtype
+        )
+        latents_std = torch.ones(1, pipe.vae.config.z_dim, 1, 1, 1, device=torch_device, dtype=pipe.vae.dtype)
+        precomputed = torch.randn(
+            1, pipe.vae.config.z_dim, num_latent_frames_per_chunk, 2, 2, device=torch_device
+        )
+
+        first_frame_latent, latents = pipe.prepare_video_latents(
+            video,
+            latents_mean=latents_mean,
+            latents_std=latents_std,
+            num_latent_frames_per_chunk=num_latent_frames_per_chunk,
+            device=torch_device,
+            latents=precomputed,
+        )
+
+        self.assertEqual(first_frame_latent.shape[2], 1)
+        torch.testing.assert_close(latents, precomputed.to(latents.dtype))
+
     @unittest.skip("Helios uses a lot of mixed precision internally, which is not suitable for this test case")
     def test_save_load_float16(self):
         pass

tests/pipelines/ltx/test_ltx_i2v_long_scheduler.py

@@ -0,0 +1,55 @@
+# Copyright 2025 The HuggingFace Team. 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.
+
+import unittest
+
+import torch
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+
+class LTXLongSchedulerStepIndexTest(unittest.TestCase):
+    def test_step_index_resync_after_skipped_steps(self):
+        sample = torch.randn(1, 4)
+        model_output = torch.randn_like(sample)
+
+        reference_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
+        reference_scheduler.set_timesteps(num_inference_steps=5, device="cpu")
+        reference = sample.clone()
+        reference = reference_scheduler.step(
+            model_output, reference_scheduler.timesteps[0], reference, return_dict=False
+        )[0]
+        reference_scheduler._step_index = None
+        reference = reference_scheduler.step(
+            model_output, reference_scheduler.timesteps[2], reference, return_dict=False
+        )[0]
+
+        stale_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
+        stale_scheduler.set_timesteps(num_inference_steps=5, device="cpu")
+        stale = sample.clone()
+        stale = stale_scheduler.step(model_output, stale_scheduler.timesteps[0], stale, return_dict=False)[0]
+        stale_wrong = stale_scheduler.step(
+            model_output, stale_scheduler.timesteps[2], stale.clone(), return_dict=False
+        )[0]
+
+        resync_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
+        resync_scheduler.set_timesteps(num_inference_steps=5, device="cpu")
+        resync = sample.clone()
+        resync = resync_scheduler.step(model_output, resync_scheduler.timesteps[0], resync, return_dict=False)[0]
+        resync_scheduler._step_index = None
+        resync = resync_scheduler.step(
+            model_output, resync_scheduler.timesteps[2], resync.clone(), return_dict=False
+        )[0]
+
+        torch.testing.assert_close(resync, reference)