190 lines
8.3 KiB
Python
190 lines
8.3 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
"""CFG Parallel Mixin for Qwen-Image 2.1.
|
||
|
|
|
||
|
|
Differences from the 2.0 mixin (``qwen_image/cfg_parallel.py``):
|
||
|
|
- The transformer owns a KV cache for the text/condition-image prefix, so the
|
||
|
|
denoise loop passes a fresh per-generation ``kv_cache`` and names the CFG
|
||
|
|
branch (``cache_branch="cond"/"uncond"``) on every forward.
|
||
|
|
- The joint sequence is ``[condition images, ..., target image]`` — the target
|
||
|
|
noise prediction is the *tail* of the output, not the head.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.distributed.cfg_parallel import (
|
||
|
|
CFGParallelMixin,
|
||
|
|
_get_cfg_world_size_or_one,
|
||
|
|
_unwrap,
|
||
|
|
_wrap,
|
||
|
|
)
|
||
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
||
|
|
get_cfg_group,
|
||
|
|
get_classifier_free_guidance_rank,
|
||
|
|
)
|
||
|
|
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class QwenImage21CFGParallelMixin(CFGParallelMixin, ProgressBarMixin):
|
||
|
|
"""Base Mixin class for the Qwen-Image 2.1 pipeline providing shared CFG methods."""
|
||
|
|
|
||
|
|
def predict_noise_maybe_with_cfg(
|
||
|
|
self,
|
||
|
|
do_true_cfg: bool,
|
||
|
|
true_cfg_scale: float,
|
||
|
|
positive_kwargs: dict[str, Any],
|
||
|
|
negative_kwargs: dict[str, Any] | None,
|
||
|
|
cfg_normalize: bool = False,
|
||
|
|
output_slice: int | None = None,
|
||
|
|
kwargs: dict[str, Any] | None = None,
|
||
|
|
) -> torch.Tensor | tuple[torch.Tensor, ...]:
|
||
|
|
# The base implementation slices `[:, :output_slice]` per branch, which
|
||
|
|
# matches the 2.0 condition-last layout. 2.1 concatenates condition
|
||
|
|
# latents *before* the target, so the target prediction is the tail of
|
||
|
|
# the joint output. Each branch must be tail-sliced *before* the CFG
|
||
|
|
# combine: positive and negative prompts tokenize to different lengths,
|
||
|
|
# so the full joint outputs do not align positionally.
|
||
|
|
def tail_slice(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
|
||
|
|
if output_slice is None:
|
||
|
|
return _wrap(pred)
|
||
|
|
return tuple(p[:, -output_slice:] for p in _wrap(pred))
|
||
|
|
|
||
|
|
if do_true_cfg:
|
||
|
|
if _get_cfg_world_size_or_one() > 1:
|
||
|
|
cfg_group = get_cfg_group()
|
||
|
|
cfg_rank = get_classifier_free_guidance_rank()
|
||
|
|
branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
|
||
|
|
local_pred = tail_slice(self.predict_noise(**branch_kwargs))
|
||
|
|
gathered = [cfg_group.all_gather(p, separate_tensors=True) for p in local_pred]
|
||
|
|
positive_noise_pred = tuple(g[0] for g in gathered)
|
||
|
|
negative_noise_pred = tuple(g[1] for g in gathered)
|
||
|
|
else:
|
||
|
|
positive_noise_pred = tail_slice(self.predict_noise(**positive_kwargs))
|
||
|
|
negative_noise_pred = tail_slice(self.predict_noise(**negative_kwargs))
|
||
|
|
return self.combine_cfg_noise(
|
||
|
|
positive_noise_pred,
|
||
|
|
negative_noise_pred,
|
||
|
|
true_cfg_scale,
|
||
|
|
cfg_normalize,
|
||
|
|
**({} if kwargs is None else {"kwargs": kwargs}),
|
||
|
|
)
|
||
|
|
|
||
|
|
pred = self.predict_noise(**positive_kwargs)
|
||
|
|
if output_slice is not None:
|
||
|
|
pred = _unwrap(tail_slice(pred))
|
||
|
|
return pred
|
||
|
|
|
||
|
|
def diffuse(
|
||
|
|
self,
|
||
|
|
prompt_embeds: torch.Tensor,
|
||
|
|
prompt_embeds_mask: torch.Tensor,
|
||
|
|
negative_prompt_embeds: torch.Tensor | None,
|
||
|
|
negative_prompt_embeds_mask: torch.Tensor | None,
|
||
|
|
img_mask: torch.Tensor,
|
||
|
|
negative_img_mask: torch.Tensor | None,
|
||
|
|
latents: torch.Tensor,
|
||
|
|
image_latents: torch.Tensor | None,
|
||
|
|
img_shapes: list,
|
||
|
|
timesteps: torch.Tensor,
|
||
|
|
do_true_cfg: bool,
|
||
|
|
true_cfg_scale: float,
|
||
|
|
attention_kwargs: dict[str, Any] | None = None,
|
||
|
|
additional_transformer_kwargs: dict[str, Any] | None = None,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""
|
||
|
|
Diffusion loop with optional classifier-free guidance.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prompt_embeds: Positive prompt embeddings
|
||
|
|
prompt_embeds_mask: Mask for positive prompt
|
||
|
|
negative_prompt_embeds: Negative prompt embeddings
|
||
|
|
negative_prompt_embeds_mask: Mask for negative prompt
|
||
|
|
img_mask: Joint-sequence image-slot mask for the positive branch
|
||
|
|
(VLM image slots plus one slot per 2x2 group of target tokens)
|
||
|
|
negative_img_mask: Same for the negative branch
|
||
|
|
latents: Noise latents to denoise (target image only, packed)
|
||
|
|
image_latents: Packed condition-image latents, prepended to the
|
||
|
|
noise latents (default: None for pure text-to-image)
|
||
|
|
img_shapes: Per-sample list of (frame, height, width) latent-token
|
||
|
|
shapes, condition images first and the target last
|
||
|
|
timesteps: Diffusion timesteps
|
||
|
|
do_true_cfg: Whether to apply CFG
|
||
|
|
true_cfg_scale: CFG scale factor
|
||
|
|
attention_kwargs: Passed through to the transformer (attn_path hint)
|
||
|
|
additional_transformer_kwargs: Extra kwargs for the transformer
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Denoised latents
|
||
|
|
"""
|
||
|
|
# Fresh cache for every generation: the pipeline is a singleton shared
|
||
|
|
# across requests, so cache state must never leak between them.
|
||
|
|
cache_enabled = getattr(self.transformer, "causal_condition", False)
|
||
|
|
kv_cache = [{} for _ in range(len(self.transformer.transformer_blocks))] if cache_enabled else None
|
||
|
|
self.scheduler.set_begin_index(0)
|
||
|
|
self.transformer.do_true_cfg = do_true_cfg
|
||
|
|
additional_transformer_kwargs = additional_transformer_kwargs or {}
|
||
|
|
|
||
|
|
with self.progress_bar(total=len(timesteps)) as pbar:
|
||
|
|
for i, t in enumerate(timesteps):
|
||
|
|
if self.interrupt:
|
||
|
|
continue
|
||
|
|
self._current_timestep = t
|
||
|
|
|
||
|
|
# Broadcast timestep to match batch size
|
||
|
|
timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
|
||
|
|
|
||
|
|
# Condition image latents go in front of the noise latents.
|
||
|
|
latent_model_input = latents
|
||
|
|
if image_latents is not None:
|
||
|
|
latent_model_input = torch.cat([image_latents, latents], dim=1)
|
||
|
|
|
||
|
|
positive_kwargs = {
|
||
|
|
"hidden_states": latent_model_input,
|
||
|
|
"timestep": timestep / 1000,
|
||
|
|
"encoder_hidden_states_mask": prompt_embeds_mask,
|
||
|
|
"encoder_hidden_states": prompt_embeds,
|
||
|
|
"img_shapes": img_shapes,
|
||
|
|
"img_mask": img_mask,
|
||
|
|
"attention_kwargs": attention_kwargs,
|
||
|
|
"kv_cache": kv_cache,
|
||
|
|
"cache_branch": "cond",
|
||
|
|
**additional_transformer_kwargs,
|
||
|
|
}
|
||
|
|
if do_true_cfg:
|
||
|
|
negative_kwargs = {
|
||
|
|
"hidden_states": latent_model_input,
|
||
|
|
"timestep": timestep / 1000,
|
||
|
|
"encoder_hidden_states_mask": negative_prompt_embeds_mask,
|
||
|
|
"encoder_hidden_states": negative_prompt_embeds,
|
||
|
|
"img_shapes": img_shapes,
|
||
|
|
"img_mask": negative_img_mask,
|
||
|
|
"attention_kwargs": attention_kwargs,
|
||
|
|
"kv_cache": kv_cache,
|
||
|
|
"cache_branch": "uncond",
|
||
|
|
**additional_transformer_kwargs,
|
||
|
|
}
|
||
|
|
else:
|
||
|
|
negative_kwargs = None
|
||
|
|
|
||
|
|
# Prefill (first step) returns the full joint sequence and
|
||
|
|
# caches the prefix; decode steps return target tokens only.
|
||
|
|
# Slicing the target tail is correct for both.
|
||
|
|
noise_pred = self.predict_noise_maybe_with_cfg(
|
||
|
|
do_true_cfg,
|
||
|
|
true_cfg_scale,
|
||
|
|
positive_kwargs,
|
||
|
|
negative_kwargs,
|
||
|
|
cfg_normalize=False,
|
||
|
|
output_slice=latents.size(1),
|
||
|
|
)
|
||
|
|
|
||
|
|
latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
|
||
|
|
|
||
|
|
pbar.update()
|
||
|
|
|
||
|
|
return latents
|