# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Base pipeline class for Diffusion models with shared CFG functionality. """ from abc import ABCMeta from typing import Any import torch from vllm.logger import init_logger from vllm.sequence import IntermediateTensors from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, get_classifier_free_guidance_world_size, ) logger = init_logger(__name__) def _wrap(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: """Normalize prediction to tuple form.""" return pred if isinstance(pred, tuple) else (pred,) def _unwrap(pred: tuple[torch.Tensor, ...]) -> torch.Tensor | tuple[torch.Tensor, ...]: """Unwrap single-element tuple to plain tensor; keep multi-element as tuple.""" return pred[0] if len(pred) == 1 else pred def _slice_pred(pred: tuple[torch.Tensor, ...], output_slice: int) -> tuple[torch.Tensor, ...]: """Slice each element along dim 1.""" return tuple(p[:, :output_slice] for p in pred) def _dispatch_branches(n_branches: int, n_ranks: int) -> list[list[int]]: """ Round-robin dispatch N branches to M ranks. Rule: branch i → rank (i % n_ranks). Examples: _dispatch_branches(3, 2) -> [[0, 2], [1]] _dispatch_branches(3, 3) -> [[0], [1], [2]] _dispatch_branches(4, 2) -> [[0, 2], [1, 3]] _dispatch_branches(4, 4) -> [[0], [1], [2], [3]] """ assignments: list[list[int]] = [[] for _ in range(n_ranks)] for i in range(n_branches): assignments[i % n_ranks].append(i) return assignments class CFGParallelMixin(metaclass=ABCMeta): """ Base Mixin class for Diffusion pipelines providing shared CFG methods. All pipelines should inherit from this class to reuse classifier-free guidance logic. CFG Parallel Architecture: When cfg_world_size > 1, each rank computes one branch (positive or negative), then all_gather exchanges results. All ranks then compute the CFG combine and scheduler step locally — no broadcast needed because the operations are deterministic. Multi-output models: Models that return tuple from predict_noise() (e.g., video + audio) should override combine_cfg_noise() to define per-element combine logic, and set self.scheduler to a composite scheduler that handles tuples. """ 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 = True, output_slice: int | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Predict noise with optional classifier-free guidance. Args: do_true_cfg: Whether to apply CFG true_cfg_scale: CFG scale factor positive_kwargs: Kwargs for positive/conditional prediction negative_kwargs: Kwargs for negative/unconditional prediction cfg_normalize: Whether to normalize CFG output (default: True) output_slice: If set, slice each output to [:, :output_slice] for image editing Returns: Predicted noise tensor or tuple of tensors. In CFG parallel mode, result is valid on ALL ranks (not just rank 0). Note: For multi-output models (e.g., video + audio where predict_noise returns a tuple), override combine_cfg_noise() for per-element CFG logic and set self.scheduler to a composite scheduler. """ if do_true_cfg: # Automatically detect CFG parallel configuration cfg_parallel_ready = get_classifier_free_guidance_world_size() > 1 if cfg_parallel_ready: cfg_group = get_cfg_group() cfg_rank = get_classifier_free_guidance_rank() # Each rank computes one branch kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs local_pred = _wrap(self.predict_noise(**kwargs)) if output_slice is not None: local_pred = _slice_pred(local_pred, output_slice) # All-gather each element, reconstruct positive/negative tuples 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) # All ranks compute combine (deterministic, same result) return self.combine_cfg_noise( positive_noise_pred, negative_noise_pred, true_cfg_scale, cfg_normalize, ) else: # Sequential CFG: compute both positive and negative positive_noise_pred = _wrap(self.predict_noise(**positive_kwargs)) negative_noise_pred = _wrap(self.predict_noise(**negative_kwargs)) if output_slice is not None: positive_noise_pred = _slice_pred(positive_noise_pred, output_slice) negative_noise_pred = _slice_pred(negative_noise_pred, output_slice) return self.combine_cfg_noise( positive_noise_pred, negative_noise_pred, true_cfg_scale, cfg_normalize, ) else: # No CFG: only compute positive/conditional prediction pred = self.predict_noise(**positive_kwargs) if output_slice is not None: pred = _unwrap(_slice_pred(_wrap(pred), output_slice)) return pred def cfg_normalize_function(self, noise_pred: torch.Tensor, comb_pred: torch.Tensor) -> torch.Tensor: """ Normalize the combined noise prediction. Args: noise_pred: positive noise prediction comb_pred: combined noise prediction after CFG Returns: Normalized noise prediction tensor """ cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) noise_pred = comb_pred * (cond_norm / noise_norm) return noise_pred def combine_cfg_noise( self, positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], true_cfg_scale: float, cfg_normalize: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Combine conditional and unconditional noise predictions with CFG. Accepts both plain tensors (backward-compatible, used by LTX2 etc.) and tuples (multi-output models). Default implementation applies the standard CFG formula to every element. Multi-output models can override this to apply different logic per element. Example override for a model returning (video_pred, audio_pred):: def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize): (video_pos, audio_pos) = positive_noise_pred (video_neg, audio_neg) = negative_noise_pred video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize) return (video_combined, audio_pos) # audio: positive only, no CFG Args: positive_noise_pred: Positive/conditional prediction(s) — Tensor or tuple negative_noise_pred: Negative/unconditional prediction(s) — Tensor or tuple true_cfg_scale: CFG scale factor cfg_normalize: Whether to normalize the combined prediction (default: False) Returns: Combined noise prediction(s) — same type as inputs """ pos_t = _wrap(positive_noise_pred) neg_t = _wrap(negative_noise_pred) results = [] for p, n in zip(pos_t, neg_t): comb = n + true_cfg_scale * (p - n) if cfg_normalize: comb = self.cfg_normalize_function(p, comb) results.append(comb) return _unwrap(tuple(results)) # ── N-branch CFG interface (for 3+ branch models) ── def predict_noise_with_multi_branch_cfg( self, do_true_cfg: bool, true_cfg_scale: float | dict[str, float], branches_kwargs: list[dict[str, Any]], cfg_normalize: bool = False, output_slice: int | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Predict noise with N-branch CFG dispatch across M GPUs. This is the multi-branch counterpart of predict_noise_maybe_with_cfg(). Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel, DreamID). Existing 2-branch models should continue using predict_noise_maybe_with_cfg(). Args: do_true_cfg: Whether to apply CFG. true_cfg_scale: CFG scale factor (passed to combine_multi_branch_cfg_noise). branches_kwargs: List of N dicts, each containing kwargs for one predict_noise() call. branches_kwargs[0] is always the positive/conditional branch. cfg_normalize: Whether to normalize (passed to combine_multi_branch_cfg_noise). output_slice: If set, slice each output to [:, :output_slice]. Returns: Combined noise prediction, identical on all ranks in CFG parallel. """ if do_true_cfg: n_branches = len(branches_kwargs) cfg_world_size = get_classifier_free_guidance_world_size() cfg_parallel_ready = cfg_world_size > 1 if cfg_parallel_ready: return self._predict_multi_branch_parallel( branches_kwargs, n_branches, cfg_world_size, true_cfg_scale, cfg_normalize, output_slice, ) else: # Sequential: run all N branches on single device preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = [] for kw in branches_kwargs: pred = _wrap(self.predict_noise(**kw)) if output_slice is not None: pred = _slice_pred(pred, output_slice) preds.append(_unwrap(pred)) return self.combine_multi_branch_cfg_noise(preds, true_cfg_scale, cfg_normalize) else: # No CFG: only compute positive/conditional prediction pred = self.predict_noise(**branches_kwargs[0]) if output_slice is not None: pred = _unwrap(_slice_pred(_wrap(pred), output_slice)) return pred def _predict_multi_branch_parallel( self, branches_kwargs: list[dict[str, Any]], n_branches: int, cfg_world_size: int, true_cfg_scale: float, cfg_normalize: bool, output_slice: int | None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Dispatch N branches across M ranks, all_gather, then combine.""" cfg_group = get_cfg_group() cfg_rank = get_classifier_free_guidance_rank() if cfg_world_size > n_branches: logger.warning_once( "cfg_parallel_size=%d > n_branches=%d, %d GPU(s) will be idle for CFG", cfg_world_size, n_branches, cfg_world_size - n_branches, ) # Assign branches to ranks via round-robin assignments = _dispatch_branches(n_branches, cfg_world_size) my_branch_ids = assignments[cfg_rank] max_per_rank = max(len(a) for a in assignments) # Run assigned branches my_preds: list[tuple[torch.Tensor, ...]] = [] for bid in my_branch_ids: pred = _wrap(self.predict_noise(**branches_kwargs[bid])) if output_slice is not None: pred = _slice_pred(pred, output_slice) my_preds.append(pred) # Idle ranks (cfg_world_size > n_branches) run a forward pass to get the output shape for all_gather. # Output shape cannot be inferred from kwargs — may be tuple, sliced, etc. if not my_preds: pred = _wrap(self.predict_noise(**branches_kwargs[0])) if output_slice is not None: pred = _slice_pred(pred, output_slice) my_preds.append(pred) # Pad to max_per_rank with zeros so all ranks have same size ref_pred = my_preds[0] while len(my_preds) < max_per_rank: my_preds.append(tuple(torch.zeros_like(t) for t in ref_pred)) # All-gather each output element separately (like predict_noise_maybe_with_cfg) # For each slot, gather across ranks; then pick valid results by owner_rank # all_slots[slot][elem_idx] = [rank0_tensor, rank1_tensor, ...] all_slots: list[list[list[torch.Tensor]]] = [] for slot in range(max_per_rank): slot_results: list[list[torch.Tensor]] = [] for p in my_preds[slot]: gathered = cfg_group.all_gather(p, separate_tensors=True) slot_results.append(gathered) all_slots.append(slot_results) # Reconstruct final_preds in branch order final_preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = [] for bid in range(n_branches): owner_rank = bid % cfg_world_size slot_idx = bid // cfg_world_size elements = tuple(all_slots[slot_idx][elem_idx][owner_rank] for elem_idx in range(len(ref_pred))) final_preds.append(_unwrap(elements)) return self.combine_multi_branch_cfg_noise(final_preds, true_cfg_scale, cfg_normalize) def combine_multi_branch_cfg_noise( self, predictions: list[torch.Tensor | tuple[torch.Tensor, ...]], true_cfg_scale: float | dict[str, float], cfg_normalize: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Combine N branch predictions. Default: standard 2-branch CFG formula. Override this method for custom multi-branch combine logic. Args: predictions: List of N predictions, where predictions[0] is always the positive/conditional branch. true_cfg_scale: CFG scale factor (float for 2-branch, dict for multi-branch). cfg_normalize: Whether to normalize the combined prediction. Returns: Combined noise prediction. """ positive = _wrap(predictions[0]) negative = _wrap(predictions[1]) results = [] for p, n in zip(positive, negative): comb = n + true_cfg_scale * (p - n) if cfg_normalize: comb = self.cfg_normalize_function(p, comb) results.append(comb) return _unwrap(tuple(results)) def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...] | IntermediateTensors: """ Forward pass through transformer to predict noise. Subclasses should override this if they need custom behavior, but the default implementation calls self.transformer. Returns: Single Tensor for standard models, or tuple of Tensors for multi-output models (e.g., video + audio). Multi-output models must also override combine_cfg_noise() and set self.scheduler to a composite scheduler that handles tuples. Non-last Pipeline Parallel stages return `IntermediateTensors` instead of final noise tensors wrapped in a tuple. """ result = self.transformer(*args, **kwargs) return result if isinstance(result, IntermediateTensors) else result[0] def diffuse( self, *args: Any, **kwargs: Any, ) -> Any: """ Diffusion loop with optional classifier-free guidance. Subclasses MUST implement this method to define the complete diffusion/denoising loop for their specific model. Typical implementation pattern (single output): ```python def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds, ...): for t in timesteps: positive_kwargs = {...} negative_kwargs = {...} noise_pred = self.predict_noise_maybe_with_cfg( do_true_cfg=do_true_cfg, true_cfg_scale=self.guidance_scale, positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, ) latents = self.scheduler_step_maybe_with_cfg( noise_pred, t, latents, do_true_cfg=do_true_cfg ) return latents ``` Multi-output models (e.g., video + audio) should: 1. Override ``predict_noise()`` to return a tuple 2. Override ``combine_cfg_noise()`` for per-element CFG logic 3. Set ``self.scheduler`` to a composite scheduler that handles tuples ```python def diffuse(self, video_latents, audio_latents, timesteps_video, timesteps_audio, ...): for t_v, t_a in zip(timesteps_video, timesteps_audio): positive_kwargs = {...} negative_kwargs = {...} # Returns tuple: (video_pred, audio_pred) video_pred, audio_pred = self.predict_noise_maybe_with_cfg( do_true_cfg=do_true_cfg, true_cfg_scale=self.guidance_scale, positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, ) # self.scheduler = VideoAudioScheduler(video_sched, audio_sched) # which accepts and returns tuples video_latents, audio_latents = self.scheduler_step_maybe_with_cfg( (video_pred, audio_pred), (t_v, t_a), (video_latents, audio_latents), do_true_cfg=do_true_cfg, ) return video_latents, audio_latents ``` """ raise NotImplementedError("Subclasses must implement diffuse") def scheduler_step( self, noise_pred: torch.Tensor | tuple[torch.Tensor, ...], t: torch.Tensor | tuple[torch.Tensor, ...], latents: torch.Tensor | tuple[torch.Tensor, ...], per_request_scheduler: Any | None = None, generator: torch.Generator | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Step the scheduler. Default implementation passes inputs directly to ``self.scheduler.step()``. For multi-output models, set ``self.scheduler`` to a composite scheduler that handles tuples (e.g., ``VideoAudioScheduler``). Args: noise_pred: Predicted noise (Tensor or tuple for multi-output) t: Current timestep (Tensor or tuple when schedulers differ per output) latents: Current latents (Tensor or tuple for multi-output) per_request_scheduler: Optional request-scoped scheduler that overrides ``self.scheduler`` for this call. This is primarily used by step-wise execution, where each request may keep scheduler state in its own runner-managed state object. Request-level execution should usually leave this as ``None`` and continue using ``self.scheduler``. generator: Optional torch Generator for reproducible sampling. When using CFG parallel, both ranks should receive generators initialized with the same seed so that non-deterministic schedulers (e.g., DDPM) produce identical results. Returns: Updated latents after scheduler step """ sched = per_request_scheduler if per_request_scheduler is not None else getattr(self, "scheduler", None) if sched is None: raise ValueError("No scheduler is available. Set self.scheduler or pass per_request_scheduler.") if not callable(getattr(sched, "step", None)): raise TypeError("per_request_scheduler must provide a callable step(...) method.") step_kwargs = dict(return_dict=False) if generator is not None: step_kwargs["generator"] = generator return sched.step(noise_pred, t, latents, **step_kwargs)[0] def scheduler_step_maybe_with_cfg( self, noise_pred: torch.Tensor | tuple[torch.Tensor, ...], t: torch.Tensor | tuple[torch.Tensor, ...], latents: torch.Tensor | tuple[torch.Tensor, ...], do_true_cfg: bool, per_request_scheduler: Any | None = None, generator: torch.Generator | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Step the scheduler with automatic CFG parallel handling. All ranks compute the scheduler step locally — no broadcast needed because predict_noise_maybe_with_cfg already ensures all ranks have identical noise_pred after all_gather + local combine. Args: noise_pred: Predicted noise (Tensor or tuple, valid on all ranks) t: Current timestep (Tensor or tuple when schedulers differ per output) latents: Current latents (Tensor or tuple) do_true_cfg: Whether CFG is enabled per_request_scheduler: Optional request-scoped scheduler that overrides ``self.scheduler`` for this call. This is mainly needed by step-wise execution, where scheduler state may be stored per request. Request-level execution should normally leave this as ``None``. generator: Optional torch Generator for reproducible sampling. When using CFG parallel, both ranks should receive generators initialized with the same seed so that non-deterministic schedulers (e.g., DDPM) produce identical results. Returns: Updated latents (identical across all CFG ranks) """ return self.scheduler_step( noise_pred, t, latents, per_request_scheduler=per_request_scheduler, generator=generator, )