diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py index 345eac107..a5ceb0cc2 100644 --- a/tests/diffusion/models/bagel/test_trajectory_recording.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index ce6c3bea1..38de0e0c5 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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