[BugFix][Bagel]: Fix vLLM-Omni as rollout bug: number of trajectory_latents count less (#3258)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -121,7 +121,8 @@ class TestTrajectoryRecording:
 
         assert trajectory_latents is not None
         assert trajectory_timesteps is not None
-        assert len(trajectory_latents) == EXPECTED_STEPS
+        # initial latent + one per denoising step
+        assert len(trajectory_latents) == EXPECTED_STEPS + 1
         assert len(trajectory_timesteps) == EXPECTED_STEPS
         # log_probs is None without a scheduler (default ODE path)
         assert trajectory_log_probs is None
@@ -167,6 +168,54 @@ class TestTrajectoryRecording:
             "Last trajectory latent should match the final output"
         )
 
+    def test_initial_latent_matches_input_noise(self, bagel_and_args):
+        """Regression: the first trajectory entry must be the initial noise (pre-denoising)."""
+        bagel, args = bagel_and_args
+        init_noise = args["packed_init_noises"].clone()
+
+        _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        assert torch.allclose(trajectory_latents[0], init_noise, atol=1e-6), (
+            "trajectory_latents[0] should be the initial noise before any denoising step"
+        )
+
+    def test_trajectory_timesteps_match_expected_schedule(self, bagel_and_args):
+        """Regression: trajectory_timesteps must record raw timesteps[i], not timesteps[i] - dts[i]."""
+        bagel, args = bagel_and_args
+
+        # Recompute the expected timestep schedule (mirrors generate_image logic)
+        ts = torch.linspace(1, 0, args["num_timesteps"])
+        shift = args.get("timestep_shift", 1.0)
+        ts = shift * ts / (1 + (shift - 1) * ts)
+        expected_timesteps = ts[:-1]  # last element is dropped
+
+        _, _, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        assert len(trajectory_timesteps) == len(expected_timesteps)
+        for i, (actual, expected) in enumerate(zip(trajectory_timesteps, expected_timesteps)):
+            assert abs(float(actual) - float(expected)) < 1e-6, (
+                f"Step {i}: trajectory_timestep={float(actual):.6f}, expected={float(expected):.6f}"
+            )
+
+    def test_latent_count_equals_timesteps_plus_one(self, bagel_and_args):
+        """Regression: len(trajectory_latents) == len(timesteps) + 1 (initial + one per step)."""
+        bagel, args = bagel_and_args
+
+        # Compute the number of denoising steps the same way generate_image does
+        ts = torch.linspace(1, 0, args["num_timesteps"])
+        shift = args.get("timestep_shift", 1.0)
+        ts = shift * ts / (1 + (shift - 1) * ts)
+        num_steps = len(ts) - 1  # timesteps = ts[:-1]
+
+        _, trajectory_latents, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        assert len(trajectory_latents) == num_steps + 1, (
+            f"Expected {num_steps + 1} latents (initial + {num_steps} steps), got {len(trajectory_latents)}"
+        )
+        assert len(trajectory_timesteps) == num_steps, (
+            f"Expected {num_steps} timesteps, got {len(trajectory_timesteps)}"
+        )
+
 
 # ---------------------------------------------------------------------------
 # Mock scheduler for log-prob tests

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1842,6 +1842,8 @@ class Bagel(nn.Module):
         # ── SP + CFG: sequential single-branch forwards ──
         use_sp = self._sp_size > 1
         if use_sp and use_cfg_text:
+            if return_trajectory_latents and len(timesteps) > 0:
+                trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
                 in_cfg_window = t > cfg_interval[0] and t <= cfg_interval[1]
@@ -1905,13 +1907,15 @@ class Bagel(nn.Module):
                     x_t = x_t - v_t.to(x_t.device) * dts[i]
                 if return_trajectory_latents:
                     trajectory_latents.append(x_t.clone())
-                    trajectory_timesteps.append(timesteps[i] - dts[i])
+                    trajectory_timesteps.append(timesteps[i])
 
             unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
             return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
 
         # ── SP without CFG: direct single-branch loop ──
         if use_sp:
+            if return_trajectory_latents and len(timesteps) > 0:
+                trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
                 v_t = self.forward_single_branch(
@@ -1938,7 +1942,7 @@ class Bagel(nn.Module):
                     x_t = x_t - v_t.to(x_t.device) * dts[i]
                 if return_trajectory_latents:
                     trajectory_latents.append(x_t.clone())
-                    trajectory_timesteps.append(timesteps[i] - dts[i])
+                    trajectory_timesteps.append(timesteps[i])
 
             unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
             return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
@@ -1999,6 +2003,9 @@ class Bagel(nn.Module):
                 "merged_cache": self._merge_naive_caches(branches_cache),
             }
 
+        if return_trajectory_latents and len(timesteps) > 0:
+            trajectory_latents.append(x_t.clone())
+
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
             if t > cfg_interval[0] and t <= cfg_interval[1]:
@@ -2036,7 +2043,7 @@ class Bagel(nn.Module):
                 x_t = x_t - v_t.to(x_t.device) * dts[i]  # velocity pointing from data to noise
             if return_trajectory_latents:
                 trajectory_latents.append(x_t.clone())
-                trajectory_timesteps.append(timesteps[i] - dts[i])
+                trajectory_timesteps.append(timesteps[i])
 
         unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
         return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
@@ -2137,6 +2144,9 @@ class Bagel(nn.Module):
         )
         _sched_kw = scheduler_kwargs or {}
 
+        if return_trajectory_latents and len(timesteps) > 0:
+            trajectory_latents.append(x_t.clone())
+
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
             use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
@@ -2194,7 +2204,7 @@ class Bagel(nn.Module):
                 x_t = x_t - v_t.to(x_t.device) * dts[i]
             if return_trajectory_latents:
                 trajectory_latents.append(x_t.clone())
-                trajectory_timesteps.append(timesteps[i] - dts[i])
+                trajectory_timesteps.append(timesteps[i])
 
         unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
         return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs