From 2cc33bb3ea107ea7c3c3e73484a8334d12462bfc Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 22 Oct 2025 17:38:49 +0800 Subject: [PATCH] Fix tensor handling in DiffusionModelRunner for multimodal outputs - Updated assertions and tensor management to correctly handle multimodal outputs in the `DiffusionModelRunner` class. - Ensured that the output tensors are processed based on the type of multimodal outputs received, improving robustness and clarity in the code. --- vllm_omni/worker/gpu_diffusion_model_runner.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py index b9d98515b..d30cf99c5 100644 --- a/vllm_omni/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/worker/gpu_diffusion_model_runner.py @@ -132,11 +132,11 @@ class DiffusionModelRunner(OmniGPUModelRunner): pooler_output: List[Optional[torch.Tensor]] = [] if isinstance(multimodal_outputs, torch.Tensor): # If model returned a single stacked tensor, split by requests - assert outputs.shape[0] == self.input_batch.num_reqs + assert multimodal_outputs.shape[0] == self.input_batch.num_reqs for i in range(self.input_batch.num_reqs): - pooler_output.append(outputs[i].detach().to("cpu").contiguous()) + pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous()) elif isinstance(multimodal_outputs, list): - for out in outputs: + for out in multimodal_outputs: pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) elif isinstance(multimodal_outputs, dict): for out in multimodal_outputs.values():