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

[Feat] DiffusionEngine Support async batch infer (#2729)

Browse Source 
Signed-off-by: Semmer <semmer@live.cn>
Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: jader <yjader@foxmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
Ting FUasukaqaq-sjaderHongsheng LiuGao Han
 authored and GitHub committed 2026-05-06 09:47:13 +12:00
1 parent 2b85474c23
commit e49fbd8a2d
32 files changed
+1951 -190
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+20 -25
View File
@@ -44,6 +44,7 @@ from vllm_omni.diffusion.utils.size_utils import (
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -864,54 +865,48 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
def denoise_step(
self,
state: "DiffusionRequestState",
input_batch: "InputBatch",
**kwargs: Any,
) -> torch.Tensor | None:
"""One denoise step: read from *state*, delegate to CFGParallelMixin.
"""One denoise step: read from *input_batch*, delegate to CFGParallelMixin.
Reuses ``predict_noise_maybe_with_cfg`` so that CFG-parallel,
sequential-CFG, and no-CFG paths are handled identically to
``diffuse()``.
"""
del kwargs
if self.interrupt:
return None
t = state.current_timestep
t = input_batch.timesteps
self._current_timestep = t
self.transformer.do_true_cfg = state.do_true_cfg
# Normalize timestep to [batch_size] tensor
if not torch.is_tensor(t):
t = torch.tensor(t, device=state.latents.device, dtype=state.latents.dtype)
self.transformer.do_true_cfg = input_batch.do_true_cfg
positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs(
latents=state.latents,
latents=input_batch.latents,
timestep=t,
guidance=state.guidance,
prompt_embeds=state.prompt_embeds,
prompt_embeds_mask=state.prompt_embeds_mask,
img_shapes=state.img_shapes,
txt_seq_lens=state.txt_seq_lens,
do_true_cfg=state.do_true_cfg,
negative_prompt_embeds=state.negative_prompt_embeds,
negative_prompt_embeds_mask=state.negative_prompt_embeds_mask,
negative_txt_seq_lens=state.negative_txt_seq_lens,
image_latents=state.sampling.image_latent,
guidance=input_batch.guidance,
prompt_embeds=input_batch.prompt_embeds,
prompt_embeds_mask=input_batch.prompt_embeds_mask,
img_shapes=input_batch.img_shapes,
txt_seq_lens=input_batch.txt_seq_lens,
do_true_cfg=input_batch.do_true_cfg,
negative_prompt_embeds=input_batch.negative_prompt_embeds,
negative_prompt_embeds_mask=input_batch.negative_prompt_embeds_mask,
negative_txt_seq_lens=input_batch.negative_txt_seq_lens,
image_latents=input_batch.image_latents,
extra_transformer_kwargs={
"attention_kwargs": self.attention_kwargs,
"return_dict": False,
},
)
true_cfg_scale = state.sampling.true_cfg_scale or 4.0
cfg_normalize = state.sampling.cfg_normalize
return self.predict_noise_maybe_with_cfg(
state.do_true_cfg,
true_cfg_scale,
input_batch.do_true_cfg,
input_batch.true_cfg_scale,
positive_kwargs,
negative_kwargs,
cfg_normalize,
input_batch.cfg_normalize,
output_slice,
)