[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:
9 files changed
+412
-22
No files matched your search
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user