Hongsheng Liu
d4096e01fe
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
415 lines
16 KiB
Python
415 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
||
import copy
|
||
import pprint
|
||
from dataclasses import asdict, dataclass, field
|
||
from typing import Any, TypeAlias
|
||
|
||
import torch
|
||
from typing_extensions import NotRequired, TypedDict
|
||
from vllm.inputs import EmbedsPrompt, PromptType, TextPrompt, TokensPrompt
|
||
from vllm.inputs.engine import TokensInput
|
||
from vllm.sampling_params import SamplingParams
|
||
|
||
from vllm_omni.lora.request import LoRARequest
|
||
|
||
DIFFUSION_QUALITY_LEVELS: tuple[str, ...] = ("lossless", "high")
|
||
|
||
|
||
class OmniTextPrompt(TextPrompt):
|
||
"""Text prompt with optional embeddings and stage payloads.
|
||
|
||
Extends TextPrompt to support prompt embeddings and additional
|
||
information payloads for direct transfer between pipeline stages.
|
||
|
||
Attributes:
|
||
prompt_embeds: Optional tensor containing prompt embeddings
|
||
additional_information: Legacy optional dictionary containing
|
||
additional information to pass along with the prompt.
|
||
model_intermediate_buffer: Optional runner-owned stage payload.
|
||
"""
|
||
|
||
negative_prompt: NotRequired[str]
|
||
# Using modalities field to differentiate between different tasks for the same pipeline
|
||
# for example Cosmos3OmniDiffusersPipeline handles t2i and t2v in the same pipeline.
|
||
modalities: NotRequired[list[str]]
|
||
prompt_embeds: NotRequired[torch.Tensor]
|
||
negative_prompt_embeds: NotRequired[torch.Tensor]
|
||
additional_information: NotRequired[dict[str, Any]]
|
||
model_intermediate_buffer: NotRequired[dict[str, Any]]
|
||
|
||
|
||
class OmniInteractionEvent(OmniTextPrompt, total=False):
|
||
"""Prompt-event payload for mid-stream interactions."""
|
||
|
||
pass
|
||
|
||
|
||
class OmniInteractionPrompt(TypedDict, total=False):
|
||
"""Mid-stream interaction payload."""
|
||
|
||
event_id: NotRequired[str]
|
||
event: NotRequired[OmniInteractionEvent]
|
||
transition_chunks: NotRequired[int]
|
||
|
||
|
||
class OmniTokensPrompt(TokensPrompt):
|
||
"""Tokens prompt with optional embeddings and stage payloads.
|
||
|
||
Extends TokensPrompt to support prompt embeddings and additional
|
||
information payloads for direct transfer between pipeline stages.
|
||
|
||
Attributes:
|
||
prompt_embeds: Optional tensor containing prompt embeddings
|
||
additional_information: Legacy optional dictionary containing
|
||
additional information to pass along with the prompt.
|
||
model_intermediate_buffer: Optional runner-owned stage payload.
|
||
"""
|
||
|
||
negative_prompt: NotRequired[str]
|
||
prompt_embeds: NotRequired[torch.Tensor]
|
||
negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]
|
||
"""The embeddings of the prompt."""
|
||
|
||
# New: optional additional information dictionary
|
||
# Values may be torch.Tensor or list
|
||
additional_information: NotRequired[dict[str, Any]]
|
||
model_intermediate_buffer: NotRequired[dict[str, Any]]
|
||
|
||
|
||
class OmniTokenInputs(TokensInput):
|
||
"""Token inputs with optional embeddings and stage payloads.
|
||
|
||
Extends TokensInput to support prompt embeddings and additional
|
||
information payloads for direct transfer between pipeline stages.
|
||
|
||
Attributes:
|
||
prompt_embeds: Optional tensor containing prompt embeddings
|
||
aligned with token IDs
|
||
additional_information: Legacy optional dictionary containing
|
||
additional information to pass along with the inputs.
|
||
model_intermediate_buffer: Optional runner-owned stage payload.
|
||
"""
|
||
|
||
# New: optional prompt embeddings aligned with token ids
|
||
negative_prompt: NotRequired[str]
|
||
prompt_embeds: NotRequired[torch.Tensor]
|
||
negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]
|
||
|
||
# New: optional additional information dictionary
|
||
# Values may be torch.Tensor or list
|
||
additional_information: NotRequired[dict[str, Any]]
|
||
model_intermediate_buffer: NotRequired[dict[str, Any]]
|
||
|
||
|
||
class OmniEmbedsPrompt(EmbedsPrompt):
|
||
"""Embeddings prompt with optional stage payloads.
|
||
|
||
Extends EmbedsPrompt to support additional information payloads
|
||
for direct transfer between pipeline stages.
|
||
|
||
Attributes:
|
||
prompt_embeds: Optional tensor containing prompt embeddings
|
||
additional_information: Legacy optional dictionary containing
|
||
additional information to pass along with the prompt.
|
||
model_intermediate_buffer: Optional runner-owned stage payload.
|
||
"""
|
||
|
||
# New: optional prompt embeddings aligned with token ids
|
||
prompt_embeds: NotRequired[torch.Tensor]
|
||
negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]
|
||
|
||
# New: optional additional information dictionary
|
||
# Values may be torch.Tensor or list
|
||
additional_information: NotRequired[dict[str, Any]]
|
||
model_intermediate_buffer: NotRequired[dict[str, Any]]
|
||
|
||
|
||
class OmniCustomPrompt(TypedDict, total=False):
|
||
"""Custom prompt type for diffusion pipelines with pre-tokenized inputs.
|
||
|
||
Allows users to pass pre-tokenized prompt IDs, attention masks, and extra
|
||
arguments directly, bypassing the tokenization stage in the pipeline.
|
||
|
||
Attributes:
|
||
prompt_ids: Pre-tokenized prompt token IDs (single or batched)
|
||
negative_prompt_ids: Pre-tokenized negative prompt token IDs
|
||
prompt_mask: Attention mask tensor for the prompt
|
||
negative_prompt_mask: Attention mask tensor for the negative prompt
|
||
extra_args: Additional pipeline-specific arguments
|
||
"""
|
||
|
||
prompt_ids: list[int] | list[list[int]]
|
||
negative_prompt_ids: list[int] | list[list[int]]
|
||
prompt_mask: torch.Tensor
|
||
negative_prompt_mask: torch.Tensor
|
||
extra_args: dict[str, Any]
|
||
|
||
|
||
# Must ensure that all additional prompt types are inherited from vLLM prompt types
|
||
# Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime.
|
||
# Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate()
|
||
OmniSingletonPrompt: TypeAlias = str | list[int] | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt
|
||
"""Omni singleton prompt type extending vLLM's SingletonPrompt."""
|
||
|
||
OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt | OmniCustomPrompt
|
||
|
||
|
||
def token_inputs_omni(
|
||
prompt_token_ids: list[int],
|
||
prompt: str | None = None,
|
||
cache_salt: str | None = None,
|
||
prompt_embeds: torch.Tensor | None = None,
|
||
additional_information: dict[str, Any] | None = None,
|
||
model_intermediate_buffer: dict[str, Any] | None = None,
|
||
) -> OmniTokenInputs:
|
||
"""Construct token inputs with optional embeddings and metadata.
|
||
|
||
Creates an OmniTokenInputs object with token IDs and optional
|
||
embeddings and additional information for pipeline stage transfer.
|
||
|
||
Args:
|
||
prompt_token_ids: List of token IDs for the prompt
|
||
prompt: Optional prompt string
|
||
cache_salt: Optional cache salt for prefix caching
|
||
prompt_embeds: Optional tensor containing prompt embeddings
|
||
additional_information: Legacy optional dictionary containing
|
||
additional information (tensors or lists).
|
||
model_intermediate_buffer: Optional runner-owned stage payload.
|
||
|
||
Returns:
|
||
OmniTokenInputs instance with the provided data
|
||
"""
|
||
inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids)
|
||
|
||
if prompt is not None:
|
||
inputs["prompt"] = prompt
|
||
if cache_salt is not None:
|
||
inputs["cache_salt"] = cache_salt
|
||
if prompt_embeds is not None:
|
||
inputs["prompt_embeds"] = prompt_embeds
|
||
if additional_information is not None:
|
||
inputs["additional_information"] = additional_information
|
||
if model_intermediate_buffer is not None:
|
||
inputs["model_intermediate_buffer"] = model_intermediate_buffer
|
||
|
||
return inputs
|
||
|
||
|
||
@dataclass
|
||
class OmniDiffusionSamplingParams:
|
||
"""
|
||
The collection of sampling parameters passed to diffusion pipelines.
|
||
|
||
This dataclass contains all information needed during the diffusion pipeline
|
||
execution, allowing methods to update specific components without needing
|
||
to manage numerous individual parameters.
|
||
"""
|
||
|
||
# Additional text-related parameters
|
||
max_sequence_length: int | None = None
|
||
prompt_template: dict[str, Any] | None = None
|
||
do_classifier_free_guidance: bool = False
|
||
output_type: str | None = None
|
||
|
||
# Request-scoped quality intent. ``None`` delegates the default behavior
|
||
# to the model; explicit levels select a model-owned quality policy.
|
||
quality: str | None = None
|
||
|
||
# Batch info
|
||
num_outputs_per_prompt: int = 1
|
||
seed: int | None = None
|
||
generator: torch.Generator | list[torch.Generator] | None = None
|
||
generator_device: str | None = None
|
||
|
||
# layered info
|
||
layers: int = 4
|
||
|
||
# cfg info
|
||
cfg_normalize: bool = False
|
||
|
||
# caption language
|
||
use_en_prompt: bool = False
|
||
|
||
# different bucket in (640, 1024) to determine the condition and output resolution
|
||
resolution: int = 640
|
||
|
||
# Tracking if embeddings are already processed
|
||
is_prompt_processed: bool = False
|
||
|
||
# Latent tensors
|
||
latents: torch.Tensor | None = None
|
||
audio_latents: torch.Tensor | None = None
|
||
raw_latent_shape: torch.Tensor | None = None
|
||
noise_pred: torch.Tensor | None = None
|
||
image_latent: torch.Tensor | None = None
|
||
|
||
# Latent dimensions
|
||
height_latents: list[int] | int | None = None
|
||
width_latents: list[int] | int | None = None
|
||
num_frames: int = 1 # Default for image models
|
||
num_frames_round_down: bool = False # Whether to round down num_frames if it's not divisible by num_gpus
|
||
|
||
# Original dimensions (before VAE scaling)
|
||
height: int | None = None
|
||
width: int | None = None
|
||
fps: float | None = None # Request/output-encoding FPS; fractional rates are supported.
|
||
frame_rate: float | None = None # Floating-point rate used by the diffusion model when it differs from `fps`.
|
||
height_not_provided: bool = False
|
||
width_not_provided: bool = False
|
||
enable_frame_interpolation: bool = False
|
||
frame_interpolation_exp: int = 1
|
||
frame_interpolation_scale: float = 1.0
|
||
frame_interpolation_model_path: str | None = None
|
||
|
||
# Timesteps
|
||
timesteps: torch.Tensor | None = None
|
||
timestep: torch.Tensor | float | int | None = None
|
||
step_index: int | None = None
|
||
boundary_ratio: float | None = None
|
||
|
||
# Scheduler parameters – ``None`` means "not explicitly set by the caller";
|
||
# each pipeline's ``forward()`` decides its own model-specific default.
|
||
num_inference_steps: int | None = None
|
||
guidance_scale: float | None = None
|
||
guidance_scale_provided: bool = False
|
||
guidance_scale_2: float | None = None
|
||
guidance_scale_2_provided: bool = False
|
||
guidance_rescale: float = 0.0
|
||
strength: float | None = None # I2I: Z-Image specific now, uses to control denoising start timestep
|
||
decode_timestep: float | list[float] | None = None
|
||
decode_noise_scale: float | list[float] | None = None
|
||
eta: float = 0.0
|
||
sigmas: list[float] | None = None
|
||
|
||
true_cfg_scale: float | None = None # qwen-image specific now
|
||
|
||
n_tokens: int | None = None
|
||
extra_step_kwargs: dict[str, Any] = field(default_factory=dict)
|
||
|
||
# [Omni] KV Cache Transfer, for bagel model now
|
||
past_key_values: Any | None = None # Injected KV Cache
|
||
kv_metadata: dict[str, Any] | None = None # Metadata for KV Cache (e.g., kv_lens, ropes)
|
||
need_kv_receive: bool = True # Flag to indicate if this request expects KV transfer
|
||
|
||
# [Omni] Multi-KV for CFG: populated by model-specific cfg_kv_collect_func
|
||
cfg_text_past_key_values: Any | None = None
|
||
cfg_img_past_key_values: Any | None = None
|
||
cfg_text_kv_metadata: dict[str, Any] | None = None
|
||
cfg_img_kv_metadata: dict[str, Any] | None = None
|
||
cfg_kv_request_ids: dict[str, str] | None = None
|
||
cfg_active_branch: str | None = None
|
||
cfg_branch_roles: list[str] | None = None
|
||
cfg_branch_past_key_values: dict[str, Any] | None = None
|
||
cfg_branch_kv_metadata: dict[str, dict[str, Any]] | None = None
|
||
|
||
# Component modules
|
||
modules: dict[str, Any] = field(default_factory=dict)
|
||
|
||
return_trajectory_latents: bool = False
|
||
return_trajectory_decoded: bool = False
|
||
trajectory_timesteps: list[torch.Tensor] | None = None
|
||
trajectory_latents: torch.Tensor | None = None
|
||
|
||
# Extra parameters that might be needed by specific pipeline implementations
|
||
extra_args: dict[str, Any] = field(default_factory=dict)
|
||
|
||
# Misc
|
||
save_output: bool = True
|
||
return_frames: bool = False
|
||
|
||
# LoRA
|
||
lora_request: LoRARequest | None = None
|
||
lora_scale: float = 1.0
|
||
|
||
# STA parameters
|
||
STA_param: list | None = None
|
||
is_cfg_negative: bool = False
|
||
mask_search_final_result_pos: list[list] | None = None
|
||
mask_search_final_result_neg: list[list] | None = None
|
||
|
||
# VSA parameters
|
||
VSA_sparsity: float = 0.0
|
||
# perf_logger: PerformanceLogger | None = None
|
||
|
||
# stage logging
|
||
# logging_info: PipelineLoggingInfo = field(default_factory=PipelineLoggingInfo)
|
||
|
||
# profile
|
||
profile: bool = False
|
||
num_profiled_timesteps: int = 8
|
||
|
||
# debugging
|
||
debug: bool = False
|
||
|
||
# results
|
||
output: torch.Tensor | None = None
|
||
|
||
def __post_init__(self) -> None:
|
||
if self.quality is not None and self.quality not in DIFFUSION_QUALITY_LEVELS:
|
||
raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {self.quality!r}")
|
||
|
||
@property
|
||
def batch_size(self):
|
||
# This class is changed to only represent a single prompt request
|
||
# Only adjust batch size for number of videos per prompt
|
||
return self.num_outputs_per_prompt
|
||
|
||
@property
|
||
def resolved_frame_rate(self) -> float | None:
|
||
"""Frame rate the model should run at, or ``None`` if the user did not provide one.
|
||
|
||
Precedence is ``frame_rate`` (the model-internal rate) over ``fps`` (the
|
||
request/output-encoding alias). Returns ``None`` when neither was set, which is
|
||
the model-agnostic "fps not provided" signal the video serving layer emits when
|
||
the request omits fps. Consumers MUST treat ``None`` as "use your own default" and
|
||
guard the read (e.g. ``req.sampling_params.resolved_frame_rate or 24.0``); never
|
||
feed it into arithmetic unguarded.
|
||
"""
|
||
if self.frame_rate is not None:
|
||
return float(self.frame_rate)
|
||
|
||
fps = self.fps
|
||
if isinstance(fps, list):
|
||
fps = fps[0] if fps else None
|
||
|
||
if fps is None:
|
||
return None
|
||
|
||
return float(fps)
|
||
|
||
def __str__(self):
|
||
return pprint.pformat(asdict(self), indent=2, width=120)
|
||
|
||
def clone(self) -> "OmniDiffusionSamplingParams":
|
||
return copy.deepcopy(self)
|
||
|
||
@classmethod
|
||
def from_params(cls, params: Any) -> "OmniDiffusionSamplingParams":
|
||
"""Normalize caller params for diffusion stages.
|
||
|
||
Existing Omni params are returned unchanged. Plain vLLM SamplingParams
|
||
copy seed and matching extra_args fields; remaining extra_args stay nested.
|
||
Unsupported types raise TypeError.
|
||
"""
|
||
if isinstance(params, cls):
|
||
return params
|
||
if isinstance(params, SamplingParams):
|
||
extra = dict(getattr(params, "extra_args", None) or {})
|
||
known = set(cls.__dataclass_fields__)
|
||
mapped = {k: extra.pop(k) for k in list(extra) if k in known}
|
||
if extra:
|
||
mapped["extra_args"] = extra
|
||
seed = getattr(params, "seed", None)
|
||
if seed is not None:
|
||
mapped.setdefault("seed", seed)
|
||
return cls(**mapped)
|
||
raise TypeError(
|
||
"Diffusion stage requires OmniDiffusionSamplingParams or vllm.SamplingParams, "
|
||
f"got {type(params).__name__!r}."
|
||
)
|
||
|
||
|
||
OmniSamplingParams: TypeAlias = SamplingParams | OmniDiffusionSamplingParams
|