diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py new file mode 100644 index 000000000..7518388d2 --- /dev/null +++ b/tests/diffusion/models/bagel/test_trajectory_recording.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for BAGEL trajectory recording in the denoising loop.""" + +import types +from dataclasses import dataclass +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.diffusion.models.bagel.bagel_transformer import ( + Bagel, + NaiveCache, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_TOKENS = 8 +HIDDEN_DIM = 16 +NUM_TIMESTEPS = 5 +# generate_image uses timesteps[:-1], so actual steps = NUM_TIMESTEPS - 1 +EXPECTED_STEPS = NUM_TIMESTEPS - 1 + + +def _make_mock_bagel(): + """Create a mock Bagel with _forward_flow returning constant velocity.""" + mock = MagicMock(spec=Bagel) + mock._sp_size = 1 + + # _forward_flow returns a small constant velocity so x_t changes each step + def fake_forward_flow(self, x_t, **kwargs): + return torch.ones_like(x_t) * 0.1 + + mock._forward_flow = types.MethodType(fake_forward_flow, mock) + # _merge_naive_caches is called in the batched CFG path + mock._merge_naive_caches = types.MethodType(lambda self, caches: NaiveCache(1), mock) + + # Bind the real generate_image to our mock + mock.generate_image = types.MethodType(Bagel.generate_image, mock) + return mock + + +def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False): + """Tensor arguments for generate_image. + + Args: + cfg: If True, enable batched CFG path (cfg_text_scale > 1.0). + """ + seq_len = num_tokens + 2 # packed_seqlens includes 2 extra tokens + base = dict( + packed_text_ids=torch.zeros(2, dtype=torch.long), + packed_text_indexes=torch.tensor([0, 1], dtype=torch.long), + packed_init_noises=torch.randn(num_tokens, hidden_dim), + packed_vae_position_ids=torch.arange(num_tokens, dtype=torch.long), + packed_vae_token_indexes=torch.arange(2, seq_len, dtype=torch.long), + packed_seqlens=torch.tensor([seq_len], dtype=torch.int), + packed_position_ids=torch.arange(seq_len, dtype=torch.long), + packed_indexes=torch.arange(seq_len, dtype=torch.long), + past_key_values=NaiveCache(1), + key_values_lens=torch.tensor([0], dtype=torch.int), + packed_key_value_indexes=torch.zeros(0, dtype=torch.long), + num_timesteps=NUM_TIMESTEPS, + timestep_shift=1.0, + cfg_text_scale=1.0, + cfg_img_scale=1.0, + ) + if cfg: + base |= dict( + cfg_text_scale=4.0, + cfg_text_packed_query_indexes=torch.arange(seq_len, dtype=torch.long), + cfg_text_packed_position_ids=torch.arange(seq_len, dtype=torch.long), + cfg_text_past_key_values=NaiveCache(1), + cfg_text_key_values_lens=torch.tensor([0], dtype=torch.int), + cfg_text_packed_key_value_indexes=torch.zeros(0, dtype=torch.long), + ) + return base + + +@pytest.fixture(params=[False, True], ids=["no_cfg", "batched_cfg"]) +def bagel_and_args(request): + """Mock Bagel instance and generate_image arguments. + + Parametrized over CFG mode so every test runs on both the no-CFG + and batched-CFG code paths. + """ + cfg = request.param + with patch( + "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size", + return_value=1, + ): + yield _make_mock_bagel(), _make_generate_args(cfg=cfg) + + +class TestTrajectoryRecording: + """Tests for trajectory latent/timestep recording in generate_image.""" + + def test_trajectory_disabled_returns_none(self, bagel_and_args): + bagel, args = bagel_and_args + + unpacked, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image( + **args, return_trajectory_latents=False + ) + + assert isinstance(unpacked, (list, tuple)) + assert len(unpacked) == 1 # one sequence + assert trajectory_latents is None + assert trajectory_timesteps is None + assert trajectory_log_probs is None + + def test_trajectory_enabled_returns_correct_count(self, bagel_and_args): + bagel, args = bagel_and_args + + _, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image( + **args, return_trajectory_latents=True + ) + + assert trajectory_latents is not None + assert trajectory_timesteps is not None + assert len(trajectory_latents) == EXPECTED_STEPS + assert len(trajectory_timesteps) == EXPECTED_STEPS + # log_probs is None without a scheduler (default ODE path) + assert trajectory_log_probs is None + + def test_trajectory_latents_shape_matches_input(self, bagel_and_args): + bagel, args = bagel_and_args + expected_shape = args["packed_init_noises"].shape + + _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True) + + for i, lat in enumerate(trajectory_latents): + assert lat.shape == expected_shape, f"Step {i}: expected {expected_shape}, got {lat.shape}" + + def test_trajectory_latents_are_distinct(self, bagel_and_args): + bagel, args = bagel_and_args + + _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True) + + for i in range(1, len(trajectory_latents)): + assert not torch.equal(trajectory_latents[i], trajectory_latents[i - 1]), ( + f"Steps {i - 1} and {i} should differ" + ) + + def test_trajectory_timesteps_are_decreasing(self, bagel_and_args): + bagel, args = bagel_and_args + + _, _, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True) + + for i in range(1, len(trajectory_timesteps)): + assert trajectory_timesteps[i] < trajectory_timesteps[i - 1], ( + f"Timestep {i} ({trajectory_timesteps[i]:.4f}) should be less than " + f"timestep {i - 1} ({trajectory_timesteps[i - 1]:.4f})" + ) + + def test_trajectory_final_latent_matches_output(self, bagel_and_args): + bagel, args = bagel_and_args + + unpacked, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True) + + # Reconstruct the full final latent from unpacked pieces + final_latent = torch.cat(unpacked, dim=0) + assert torch.allclose(trajectory_latents[-1], final_latent, atol=1e-6), ( + "Last trajectory latent should match the final output" + ) + + +# --------------------------------------------------------------------------- +# Mock scheduler for log-prob tests +# --------------------------------------------------------------------------- + + +@dataclass +class _MockStepOutput: + prev_sample: torch.Tensor + log_prob: torch.Tensor + + +class _MockScheduler: + """Minimal scheduler: Euler step + constant log-prob per step.""" + + def step(self, model_output, sigma, sample, dt, **kwargs): + prev_sample = sample - model_output * dt + log_prob = torch.tensor(-1.0) + return _MockStepOutput(prev_sample=prev_sample, log_prob=log_prob) + + +class TestTrajectoryLogProbs: + """Tests for log-prob recording when a scheduler is provided.""" + + @pytest.fixture() + def bagel_scheduler_args(self): + with patch( + "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size", + return_value=1, + ): + yield _make_mock_bagel(), _make_generate_args(), _MockScheduler() + + def test_log_probs_recorded_with_scheduler(self, bagel_scheduler_args): + bagel, args, scheduler = bagel_scheduler_args + + _, _, _, trajectory_log_probs = bagel.generate_image( + **args, return_trajectory_latents=True, scheduler=scheduler + ) + + assert trajectory_log_probs is not None + assert len(trajectory_log_probs) == EXPECTED_STEPS + + def test_log_probs_are_finite(self, bagel_scheduler_args): + bagel, args, scheduler = bagel_scheduler_args + + _, _, _, trajectory_log_probs = bagel.generate_image( + **args, return_trajectory_latents=True, scheduler=scheduler + ) + + for i, lp in enumerate(trajectory_log_probs): + assert torch.isfinite(lp).all(), f"Step {i}: log_prob is not finite" + + def test_log_probs_none_without_scheduler(self, bagel_scheduler_args): + bagel, args, _ = bagel_scheduler_args + + _, _, _, trajectory_log_probs = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None) + + assert trajectory_log_probs is None + + def test_scheduler_updates_latents(self, bagel_scheduler_args): + """Verify the scheduler's prev_sample is used (not the raw Euler step).""" + bagel, args, scheduler = bagel_scheduler_args + + _, traj_with_sched, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=scheduler) + _, traj_without, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None) + + # Mock scheduler does the same Euler step, so latents should match + for i in range(len(traj_with_sched)): + assert torch.allclose(traj_with_sched[i], traj_without[i], atol=1e-5), ( + f"Step {i}: scheduler and ODE paths should produce same latents" + ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 3071fd9d5..56a891aa5 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -9,6 +9,7 @@ from dataclasses import dataclass, field, fields from typing import TYPE_CHECKING, Any import torch +from PIL import Image from pydantic import model_validator from typing_extensions import Self from vllm.config.utils import config @@ -701,10 +702,12 @@ class DiffusionOutput: Final output (after pipeline completion) """ - output: torch.Tensor | None = None - trajectory_timesteps: list[torch.Tensor] | None = None - trajectory_latents: torch.Tensor | None = None - trajectory_decoded: list[torch.Tensor] | None = None + # Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm + output: torch.Tensor | dict | None = None + trajectory_timesteps: torch.Tensor | dict | None = None + trajectory_latents: torch.Tensor | dict | None = None + trajectory_log_probs: torch.Tensor | dict | None = None + trajectory_decoded: list[Image.Image] | None = None error: str | None = None aborted: bool = False abort_message: str | None = None diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 784da6175..8d3c02b7a 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -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, diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 9aafc1cf1..6a96533fd 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -78,13 +78,29 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor: return tensor +def _pack_tensor_if_large(val: torch.Tensor) -> torch.Tensor | dict: + """Replace a tensor with an SHM handle if it exceeds the threshold.""" + if val.nelement() * val.element_size() > _SHM_TENSOR_THRESHOLD: + return _tensor_to_shm(val) + return val + + +def _unpack_if_shm_handle(val: object) -> object: + """Reconstruct a tensor from an SHM handle dict, or return as-is.""" + if isinstance(val, dict) and val.get("__tensor_shm__"): + return _tensor_from_shm(val) + return val + + def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: if output.output is not None and isinstance(output.output, torch.Tensor): - if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD: - output.output = _tensor_to_shm(output.output) + output.output = _pack_tensor_if_large(output.output) if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor): - if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD: - output.trajectory_latents = _tensor_to_shm(output.trajectory_latents) + output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents) + if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor): + output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps) + if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor): + output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs) return output @@ -104,10 +120,10 @@ def pack_diffusion_output_shm(output: object) -> object: def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: - if isinstance(output.output, dict) and output.output.get("__tensor_shm__"): - output.output = _tensor_from_shm(output.output) - if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"): - output.trajectory_latents = _tensor_from_shm(output.trajectory_latents) + output.output = _unpack_if_shm_handle(output.output) + output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents) + output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps) + output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs) return output diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index d32a6d8ac..a04ded376 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -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( diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 84f177e01..2c72d9890 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -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]: diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 0a5fd3590..bcc3bef15 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -14,8 +14,10 @@ from multiprocessing.process import BaseProcess from typing import TYPE_CHECKING, Any import msgspec +import torch import zmq import zmq.asyncio +from PIL import Image from vllm.logger import init_logger from vllm.transformers_utils.config import get_hf_file_to_dict from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx @@ -174,8 +176,13 @@ class StageDiffusionProc: merged_mm: dict[str, Any] = {} merged_metrics: dict[str, Any] = {} merged_durations: dict[str, float] = {} + merged_custom: dict[str, Any] = {} peak_mem = 0.0 latents = None + trajectory_latents: list[torch.Tensor] | None = None + trajectory_timesteps: list[torch.Tensor] | None = None + trajectory_log_probs: torch.Tensor | None = None + trajectory_decoded: list[Image.Image] | None = None final_output_type = "image" for r in results: @@ -183,9 +190,18 @@ class StageDiffusionProc: merged_mm.update(r._multimodal_output) merged_metrics.update(r.metrics) merged_durations.update(r.stage_durations) + merged_custom.update(r._custom_output) peak_mem = max(peak_mem, r.peak_memory_mb) if latents is None and r.latents is not None: latents = r.latents + if trajectory_latents is None: + trajectory_latents = r.trajectory_latents + if trajectory_timesteps is None: + trajectory_timesteps = r.trajectory_timesteps + if trajectory_log_probs is None: + trajectory_log_probs = r.trajectory_log_probs + if trajectory_decoded is None: + trajectory_decoded = r.trajectory_decoded if r.final_output_type != "image": final_output_type = r.final_output_type @@ -195,6 +211,11 @@ class StageDiffusionProc: prompt=prompts[0] if len(prompts) == 1 else None, metrics=merged_metrics, latents=latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, + custom_output=merged_custom or None, multimodal_output=merged_mm or None, final_output_type=final_output_type, stage_durations=merged_durations, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 96df0591e..1a7ffc4a5 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -282,6 +282,11 @@ class OmniBase: final_output_type=stage_meta["final_output_type"], request_output=engine_outputs, images=images, + trajectory_latents=getattr(engine_outputs, "trajectory_latents", None), + trajectory_timesteps=getattr(engine_outputs, "trajectory_timesteps", None), + trajectory_log_probs=getattr(engine_outputs, "trajectory_log_probs", None), + trajectory_decoded=getattr(engine_outputs, "trajectory_decoded", None), + _custom_output=getattr(engine_outputs, "_custom_output", {}), stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, ) diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index ca3ba271a..4a775356e 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -58,6 +58,10 @@ class OmniRequestOutput: images: list[Image.Image] = field(default_factory=list) prompt: OmniPromptType | None = None latents: torch.Tensor | None = None + trajectory_latents: torch.Tensor | None = None + trajectory_timesteps: torch.Tensor | None = None + trajectory_log_probs: torch.Tensor | None = None + trajectory_decoded: list | None = None metrics: dict[str, Any] = field(default_factory=dict) _multimodal_output: dict[str, Any] = field(default_factory=dict) _custom_output: dict[str, Any] = field(default_factory=dict) @@ -101,6 +105,10 @@ class OmniRequestOutput: prompt: OmniPromptType | None = None, metrics: dict[str, Any] | None = None, latents: torch.Tensor | None = None, + trajectory_latents: torch.Tensor | None = None, + trajectory_timesteps: torch.Tensor | None = None, + trajectory_log_probs: torch.Tensor | None = None, + trajectory_decoded: list | None = None, multimodal_output: dict[str, Any] | None = None, custom_output: dict[str, Any] | None = None, final_output_type: str = "image", @@ -129,6 +137,10 @@ class OmniRequestOutput: images=images, prompt=prompt, latents=latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, metrics=metrics or {}, _multimodal_output=multimodal_output or {}, _custom_output=custom_output or {},