[Feature] Add trajectory recording to BAGEL denoising loop (#2483)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Signed-off-by: Zhengyuan Su (苏政渊) <su.zhengyuan@u.nus.edu>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Samit <285365963@qq.com>
This commit is contained in:
authored and GitHub committed 2026-04-08 10:43:01 +08:00
1 parent 6433847249
commit cb6a8739c1
9 files changed
+412 -22

No files matched your search

+8
View File
@@ -209,6 +209,10 @@ class DiffusionEngine:
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
trajectory_latents=output.trajectory_latents,
trajectory_timesteps=output.trajectory_timesteps,
trajectory_log_probs=output.trajectory_log_probs,
trajectory_decoded=output.trajectory_decoded,
custom_output=output.custom_output or {},
multimodal_output=mm_output,
stage_durations=output.stage_durations,
@@ -267,6 +271,10 @@ class DiffusionEngine:
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
trajectory_latents=output.trajectory_latents,
trajectory_timesteps=output.trajectory_timesteps,
trajectory_log_probs=output.trajectory_log_probs,
trajectory_decoded=output.trajectory_decoded,
custom_output=output.custom_output or {},
multimodal_output=mm_output,
stage_durations=output.stage_durations,