anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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:
Zhengyuan Su (苏政渊)ClaudeSamit
 authored and GitHub committed 2026-04-08 10:43:01 +08:00
1 parent 6433847249
commit cb6a8739c1
9 files changed
+412 -22
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/bagel/bagel_transformer.py
+69 -8
View File
@@ -1655,6 +1655,9 @@ class Bagel(nn.Module):
cfg_img_past_key_values: NaiveCache | None = None,
cfg_img_key_values_lens: torch.IntTensor | None = None,
cfg_img_packed_key_value_indexes: torch.LongTensor | None = None,
return_trajectory_latents: bool = False,
scheduler: object | None = None,
scheduler_kwargs: dict | None = None,
):
x_t = packed_init_noises
@@ -1663,6 +1666,14 @@ class Bagel(nn.Module):
dts = timesteps[:-1] - timesteps[1:]
timesteps = timesteps[:-1]
# Optional trajectory recording for RL rollout data collection
trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None
trajectory_timesteps: list[torch.Tensor] | None = [] if return_trajectory_latents else None
trajectory_log_probs: list[torch.Tensor] | None = (
[] if (return_trajectory_latents and scheduler is not None) else None
)
_sched_kw = scheduler_kwargs or {}
use_cfg_text = cfg_text_scale > 1.0
use_cfg_img = cfg_img_scale > 1.0
@@ -1699,6 +1710,9 @@ class Bagel(nn.Module):
cfg_img_past_key_values=cfg_img_past_key_values,
cfg_img_key_values_lens=cfg_img_key_values_lens,
cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes,
return_trajectory_latents=return_trajectory_latents,
scheduler=scheduler,
scheduler_kwargs=scheduler_kwargs,
)
# ── SP + CFG: sequential single-branch forwards ──
@@ -1758,10 +1772,19 @@ class Bagel(nn.Module):
cfg_renorm_min,
)
x_t = x_t - v_t.to(x_t.device) * dts[i]
if scheduler is not None:
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
x_t = out.prev_sample
if trajectory_log_probs is not None and out.log_prob is not None:
trajectory_log_probs.append(out.log_prob)
else:
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])
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
return unpacked_latent
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
# ── SP without CFG: direct single-branch loop ──
if use_sp:
@@ -1781,10 +1804,20 @@ class Bagel(nn.Module):
past_key_values=past_key_values,
packed_key_value_indexes=packed_key_value_indexes,
)
x_t = x_t - v_t.to(x_t.device) * dts[i]
if scheduler is not None:
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
x_t = out.prev_sample
out_log_prob = getattr(out, "log_prob", None)
if trajectory_log_probs is not None and out_log_prob is not None:
trajectory_log_probs.append(out_log_prob)
else:
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])
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
return unpacked_latent
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
# ── Batched CFG mode (cfg_parallel_size=1, no SP) ──
cfg_batched = None
@@ -1870,10 +1903,19 @@ class Bagel(nn.Module):
cfg_batched=cfg_batched,
)
x_t = x_t - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise
if scheduler is not None:
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
x_t = out.prev_sample
if trajectory_log_probs is not None and out.log_prob is not None:
trajectory_log_probs.append(out.log_prob)
else:
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])
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
return unpacked_latent
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
def _generate_image_parallel(
self,
@@ -1905,6 +1947,9 @@ class Bagel(nn.Module):
cfg_img_past_key_values: NaiveCache | None,
cfg_img_key_values_lens: torch.IntTensor | None,
cfg_img_packed_key_value_indexes: torch.LongTensor | None,
return_trajectory_latents: bool = False,
scheduler: object | None = None,
scheduler_kwargs: dict | None = None,
):
"""CFG parallel denoising loop: each rank computes one CFG branch.
@@ -1961,6 +2006,13 @@ class Bagel(nn.Module):
else:
raise RuntimeError(f"Unexpected cfg_rank={cfg_rank} for Bagel 3-branch CFG parallel")
trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None
trajectory_timesteps: list[torch.Tensor] | None = [] if return_trajectory_latents else None
trajectory_log_probs: list[torch.Tensor] | None = (
[] if (return_trajectory_latents and scheduler is not None) else None
)
_sched_kw = scheduler_kwargs or {}
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
@@ -2009,10 +2061,19 @@ class Bagel(nn.Module):
packed_key_value_indexes=packed_key_value_indexes,
)
x_t = x_t - v_t.to(x_t.device) * dts[i]
if scheduler is not None:
out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw)
x_t = out.prev_sample
if trajectory_log_probs is not None and out.log_prob is not None:
trajectory_log_probs.append(out.log_prob)
else:
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])
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
return unpacked_latent
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
@staticmethod
def _combine_cfg(
vllm_omni/diffusion/models/bagel/pipeline_bagel.py
+30 -2
View File
@@ -159,6 +159,9 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
self.od_config = od_config
self.device = get_local_device()
self._scheduler: object | None = None
self._scheduler_kwargs: dict = {}
model = od_config.model
local_files_only = os.path.exists(model)
if local_files_only:
@@ -630,7 +633,7 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
enabled=self.device.type != "cpu",
dtype=self.od_config.dtype,
):
latents = self.bagel.generate_image(
latents, trajectory_latents, trajectory_timesteps, trajectory_log_probs = self.bagel.generate_image(
past_key_values=gen_context["past_key_values"],
cfg_text_past_key_values=cfg_text_context["past_key_values"],
cfg_img_past_key_values=cfg_img_context["past_key_values"],
@@ -650,11 +653,36 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
cfg_img_packed_query_indexes=generation_input_cfg_img["cfg_packed_query_indexes"],
cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"],
cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"],
return_trajectory_latents=req.sampling_params.return_trajectory_latents,
scheduler=self._scheduler,
scheduler_kwargs=self._scheduler_kwargs,
)
img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
# Build trajectory output when requested
trajectory_latents_stacked: torch.Tensor | None = None
trajectory_timesteps_stacked: torch.Tensor | None = None
trajectory_decoded: list[Image.Image] | None = None
if trajectory_latents:
trajectory_latents_stacked = torch.stack(trajectory_latents)
trajectory_timesteps_stacked = torch.stack(trajectory_timesteps)
if req.sampling_params.return_trajectory_decoded:
trajectory_decoded = [
self._decode_image_from_latent(self.bagel, self.vae, lat, image_shape) for lat in trajectory_latents
]
trajectory_log_probs_stacked: torch.Tensor | None = None
if trajectory_log_probs:
trajectory_log_probs_stacked = torch.stack(trajectory_log_probs)
return DiffusionOutput(
output=img, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
output=img,
trajectory_latents=trajectory_latents_stacked,
trajectory_timesteps=trajectory_timesteps_stacked,
trajectory_log_probs=trajectory_log_probs_stacked,
trajectory_decoded=trajectory_decoded,
stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: