From 28467dc19a8c453105ce30c8e97537d2d265dd00 Mon Sep 17 00:00:00 2001 From: ANHDY <49194355+ANHDY@users.noreply.github.com> Date: Fri, 22 May 2026 15:23:03 +0800 Subject: [PATCH] [feature]: support Hidream-I1-Full model (#2572) Signed-off-by: wangmengdi <160612767@qq.com> Co-authored-by: wangmengdi --- docs/models/supported_models.md | 2 + docs/user_guide/diffusion_features.md | 1 + .../offline_inference/text_to_image/README.md | 20 + .../text_to_image/text_to_image.py | 7 + pyproject.toml | 2 + recipes/HiDream/HiDream-I1-Image.md | 106 ++ .../test_hidream_i1_full_expansion.py | 69 ++ vllm_omni/diffusion/data.py | 5 +- .../models/hidream_image/__init__.py | 17 + .../hidream_image_transformer.py | 1042 +++++++++++++++++ .../pipeline_hidream_i1_image.py | 1027 ++++++++++++++++ .../hidream_image/pipeline_hidream_image.py | 1027 ++++++++++++++++ vllm_omni/diffusion/registry.py | 6 + vllm_omni/engine/async_omni_engine.py | 4 + vllm_omni/entrypoints/cli/serve.py | 10 + 15 files changed, 3344 insertions(+), 1 deletion(-) create mode 100644 recipes/HiDream/HiDream-I1-Image.md create mode 100644 tests/e2e/online_serving/test_hidream_i1_full_expansion.py create mode 100644 vllm_omni/diffusion/models/hidream_image/__init__.py create mode 100644 vllm_omni/diffusion/models/hidream_image/hidream_image_transformer.py create mode 100644 vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py create mode 100644 vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f30f34758..b0637a5d7 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -71,4 +71,6 @@ th { | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | + ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend. diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 61694ac9d..2a296a059 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -115,6 +115,7 @@ The following tables show which models support each feature: | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | | **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index c71773972..02538b40f 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -35,6 +35,7 @@ This folder provides several entrypoints for experimenting with text-to-image di | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 | | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) | | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3) | 160 | +| `HiDream-I1-Full` | 1024 x 1024 | 63.7 | 57.7 | !!! info *Peak VRAM: based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU. @@ -94,6 +95,7 @@ python text_to_image.py \ | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights | | `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.| | `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.| +| `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). | **NextStep-1.1 specific arguments:** @@ -177,6 +179,24 @@ python examples/offline_inference/text_to_image/text_to_image.py \ --output flux2-dev.png ``` + +### HiDream-I1-Full Models + +The `--auxiliary-text-encoder` parameter is required when running HiDream‑I1‑Full: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model HiDream-ai/HiDream-I1-Full \ + --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \ + --seed 42 \ + --guidance-scale 5.0 \ + --tensor-parallel-size 1 \ + --num-images-per-prompt 1 \ + --num-inference-steps 50 \ + --auxiliary-text-encoder meta-llama/Meta-Llama-3.1-8B-Instruct \ + --output /output.png +``` + ### Batch Requests (Multiple Prompts) You can pass multiple prompts in a single `generate` call. diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index c0fd337bd..443b21acc 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -300,6 +300,12 @@ def parse_args() -> argparse.Namespace: default=None, help=("Custom system prompt. Used when --use-system-prompt is custom. "), ) + parser.add_argument( + "--auxiliary-text-encoder", + type=str, + default=None, + help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).", + ) current_omni_platform.pre_register_and_update(parser) from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults @@ -392,6 +398,7 @@ def main(): "profiler_config": args.profiler_config, "init_timeout": args.init_timeout, "stage_init_timeout": args.stage_init_timeout, + "auxiliary_text_encoder": args.auxiliary_text_encoder, **lora_args, **quant_kwargs, } diff --git a/pyproject.toml b/pyproject.toml index bb9cdc688..e2c9242b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -241,6 +241,8 @@ extend-ignore-identifiers-re = [ ".*MoBA", ".*temperal_downsample", ".*_nd", + ".*nd$", + ".*EmbedND$", ".*nothink.*", ".*NOTHINK.*", ".*nin.*", diff --git a/recipes/HiDream/HiDream-I1-Image.md b/recipes/HiDream/HiDream-I1-Image.md new file mode 100644 index 000000000..958d03d6d --- /dev/null +++ b/recipes/HiDream/HiDream-I1-Image.md @@ -0,0 +1,106 @@ +# HiDream-I1-Full + +> Text-to-image serving and benchmark + +## Summary + +- Vendor: HiDream.ai +- Model: `HiDream-ai/HiDream-I1-Full` +- Task: Text-to-image generation +- Mode: Online serving with the OpenAI-compatible API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a known-good starting point for serving +`HiDream-ai/HiDream-I1-Full` on a single 80 GB A800, validate the normal online-serving path. + +## References + +- Model: +- Upstream or canonical docs: + [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md) +- Related example under `examples/`: + [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md) + + +## Hardware Support + +This recipe currently documents one CUDA GPU serving configuration. + +## GPU + +### 1x A800 80GB + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +Start the baseline server: + +```bash +export MODEL_NAME_OR_PATH=HiDream-ai/HiDream-I1-Full +vllm serve ${MODEL_NAME_OR_PATH} \ + --omni \ + --port 8092 \ + --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \ + --tensor-parallel-size 1 \ + --vae_use_slicing \ + --vae_use_tiling +``` + + +You can also use the example launcher and pass the extra flags through: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model HiDream-ai/HiDream-I1-Full \ + --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \ + --seed 42 \ + --guidance-scale 5.0 \ + --tensor-parallel-size 1 \ + --num-images-per-prompt 1 \ + --num-inference-steps 50 \ + --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \ + --output /workspace/mnt/cmss-wmd/result/output.png + +``` + +#### Verification + +Run the existing client example after the server is ready: + +```bash +python examples/online_serving/text_to_image/openai_chat_client.py \ + --server http://localhost:8092 \ + --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \ + --output /tmp/hidream_i1_recipe.png +``` + +For a direct API smoke test: + +```bash +curl -s http://localhost:8092/v1/chat/completions -H "Content-Type: application/json" -d '{ + "messages": [ + {"role": "user", "content": "The setting sun of late autumn dyes the riverside with a warm orange hue"} + ], + "extra_body": { + "num_inference_steps": 50, + "seed": 42, + "height": 1024, + "width": 1024 + } + }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output_online.png +``` + + + +#### Notes + +- Key flags: `--auxiliary-text-encoder` designates the path for the auxiliary text model meta-llama/Llama-3.1-8B-Instruct. For HiDream-I1-Full, unspecified use defaults meta-llama/Llama-3.1-8B-Instruct (downloaded from official Hugging Face), and custom paths are supported. diff --git a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py new file mode 100644 index 000000000..f84cc08ad --- /dev/null +++ b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py @@ -0,0 +1,69 @@ +""" +End-to-end diffusion coverage for HiDream-I1-Full in online serving mode. + +This test verifies that HiDream-I1-Full can be launched, +accepts text-to-image requests through the OpenAI-compatible API, and returns +valid generated images with the requested resolution. + +assert_diffusion_response validates successful generation and the expected +image resolution. +""" + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +MODEL = "HiDream-ai/HiDream-I1-Full" +PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography." +NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark" + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _get_hidream_i1_image_feature_cases(model: str): + """Return HiDream-I1-Full diffusion feature cases.""" + + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--auxiliary-text-encoder", + "meta-llama/Meta-Llama-3.1-8B-Instruct", + ], + ), + id="default", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.parametrize( + "omni_server", + _get_hidream_i1_image_feature_cases(MODEL), + indirect=True, +) +def test_hidream_i1_image( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, +): + """Validate HiDream-I1-Full online serving with CPU offload.""" + + messages = dummy_messages_from_mix_data(content_text=PROMPT) + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "negative_prompt": NEGATIVE_PROMPT, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9841aa2e3..a6fe1e4e9 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Any import diffusers import torch from PIL import Image -from pydantic import model_validator +from pydantic import Field, model_validator from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger @@ -593,6 +593,9 @@ class OmniDiffusionConfig: # Maximum number of sequences to generate in a batch max_num_seqs: int = 1 + # Supplementary model specific parameters + extras: dict[str, Any] = Field(default_factory=dict) + @property def is_moe(self) -> bool: num_experts = self.tf_model_config.get("num_experts", None) diff --git a/vllm_omni/diffusion/models/hidream_image/__init__.py b/vllm_omni/diffusion/models/hidream_image/__init__.py new file mode 100644 index 000000000..d8c52afdd --- /dev/null +++ b/vllm_omni/diffusion/models/hidream_image/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HiDream Image diffusion model components.""" + +from vllm_omni.diffusion.models.hidream_image.hidream_image_transformer import ( + HiDreamImageTransformer2DModel, +) +from vllm_omni.diffusion.models.hidream_image.pipeline_hidream_image import ( + HiDreamImagePipeline, + get_hidream_image_post_process_func, +) + +__all__ = [ + "HiDreamImageTransformer2DModel", + "HiDreamImagePipeline", + "get_hidream_image_post_process_func", +] diff --git a/vllm_omni/diffusion/models/hidream_image/hidream_image_transformer.py b/vllm_omni/diffusion/models/hidream_image/hidream_image_transformer.py new file mode 100644 index 000000000..5c12df14e --- /dev/null +++ b/vllm_omni/diffusion/models/hidream_image/hidream_image_transformer.py @@ -0,0 +1,1042 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.utils import deprecate, logging +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +logger = logging.get_logger(__name__) + + +class HiDreamImageFeedForwardSwiGLU(nn.Module): + def __init__( + self, + dim: int, + hidden_dim: int, + multiple_of: int = 256, + ffn_dim_multiplier: float | None = None, + ): + super().__init__() + hidden_dim = int(2 * hidden_dim / 3) + # custom dim factor multiplier + if ffn_dim_multiplier is not None: + hidden_dim = int(ffn_dim_multiplier * hidden_dim) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(torch.nn.functional.silu(self.w1(x)) * self.w3(x)) + + +class HiDreamImagePooledEmbed(nn.Module): + def __init__(self, text_emb_dim, hidden_size): + super().__init__() + self.pooled_embedder = TimestepEmbedding(in_channels=text_emb_dim, time_embed_dim=hidden_size) + + def forward(self, pooled_embed: torch.Tensor) -> torch.Tensor: + # + return self.pooled_embedder(pooled_embed.to(dtype=torch.bfloat16)) + + +class HiDreamImageTimestepEmbed(nn.Module): + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.time_proj = Timesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=frequency_embedding_size, time_embed_dim=hidden_size) + + def forward(self, timesteps: torch.Tensor, wdtype: torch.dtype | None = None) -> torch.Tensor: + t_emb = self.time_proj(timesteps).to(dtype=torch.bfloat16) + t_emb = self.timestep_embedder(t_emb) + return t_emb + + +class HiDreamImageOutEmbed(nn.Module): + def __init__(self, hidden_size, patch_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + shift, scale = self.adaLN_modulation(temb).chunk(2, dim=1) + hidden_states = self.norm_final(hidden_states) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + hidden_states = self.linear(hidden_states) + return hidden_states + + +class HiDreamImagePatchEmbed(nn.Module): + def __init__( + self, + patch_size=2, + in_channels=4, + out_channels=1024, + ): + super().__init__() + self.patch_size = patch_size + self.out_channels = out_channels + self.proj = nn.Linear(in_channels * patch_size * patch_size, out_channels, bias=True) + + def forward(self, latent) -> torch.Tensor: + latent = self.proj(latent.to(dtype=torch.bfloat16)) + return latent + + +def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: + assert dim % 2 == 0, "The dimension must be even." + + is_mps = pos.device.type == "mps" + is_npu = pos.device.type == "npu" + + dtype = torch.float32 if (is_mps or is_npu) else torch.float64 + + scale = torch.arange(0, dim, 2, dtype=dtype, device=pos.device) / dim + omega = 1.0 / (theta**scale) + + batch_size, seq_length = pos.shape + out = torch.einsum("...n,d->...nd", pos, omega) + cos_out = torch.cos(out) + sin_out = torch.sin(out) + + stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1) + out = stacked_out.view(batch_size, -1, dim // 2, 2, 2) + return out.float() + + +class HiDreamImageEmbedND(nn.Module): + def __init__(self, theta: int, axes_dim: list[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + n_axes = ids.shape[-1] + emb = torch.cat( + [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], + dim=-3, + ) + return emb.unsqueeze(2) + + +class DistributedRMSNorm(nn.Module): + """RMSNorm that computes global RMS across tensor parallel ranks.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + + input_dtype = x.dtype + x_float = x.float() + local_sum_sq = (x_float**2).sum(dim=-1, keepdim=True) + local_count = x.shape[-1] + + if tp_size > 1: + global_sum_sq = local_sum_sq.clone() + tensor_model_parallel_all_reduce(global_sum_sq) + global_count = local_count * tp_size + else: + global_sum_sq = local_sum_sq + global_count = local_count + + rms = torch.sqrt(global_sum_sq / global_count + self.eps) + + output = (x_float / rms) * self.weight.float() + return output.to(input_dtype) + + +class HiDreamAttention(nn.Module): + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + upcast_attention: bool = False, + upcast_softmax: bool = False, + scale_qk: bool = True, + eps: float = 1e-5, + processor=None, + out_dim: int = None, + single: bool = False, + quant_config: "QuantizationConfig | None" = None, + ): + super().__init__() + self.head_dim = dim_head + + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.query_dim = query_dim + self.upcast_attention = upcast_attention + self.upcast_softmax = upcast_softmax + self.out_dim = out_dim if out_dim is not None else query_dim + + self.scale_qk = scale_qk + self.scale = dim_head**-0.5 if self.scale_qk else 1.0 + + self.heads = out_dim // dim_head if out_dim is not None else heads + self.sliceable_head_dim = heads + self.single = single + + # Fused QKV projection using vLLM's optimized layer + self.to_qkv = QKVParallelLinear( + hidden_size=query_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + quant_config=quant_config, + prefix="to_qkv", + ) + self.query_num_heads = self.to_qkv.num_heads + self.kv_num_heads = self.to_qkv.num_kv_heads + + self.to_out = RowParallelLinear( + self.inner_dim, + self.out_dim, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix="to_out", + ) + + self.tp_inner_dim = self.query_num_heads * dim_head + if get_tensor_model_parallel_world_size() > 1: + self.q_rms_norm = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + self.k_rms_norm = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + else: + self.q_rms_norm = nn.RMSNorm(self.tp_inner_dim, eps) + self.k_rms_norm = nn.RMSNorm(self.tp_inner_dim, eps) + + if not single: + self.to_qkv_t = QKVParallelLinear( + hidden_size=query_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + quant_config=quant_config, + prefix="to_qkv_t", + ) + self.to_out_t = RowParallelLinear( + self.inner_dim, + self.out_dim, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix="to_out_t", + ) + if get_tensor_model_parallel_world_size() > 1: + self.q_rms_norm_t = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + self.k_rms_norm_t = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + else: + self.q_rms_norm_t = nn.RMSNorm(self.tp_inner_dim, eps) + self.k_rms_norm_t = nn.RMSNorm(self.tp_inner_dim, eps) + self.rope = RotaryEmbedding(is_neox_style=False) + + def forward( + self, + hidden_states: torch.Tensor, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + dtype = hidden_states.dtype + batch_size = hidden_states.shape[0] + + qkv, _ = self.to_qkv(hidden_states) + q_size = self.to_qkv.num_heads * self.head_dim + kv_size = self.to_qkv.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + + query_i = self.q_rms_norm(query).to(dtype=dtype) + key_i = self.k_rms_norm(key).to(dtype=dtype) + value_i = value + + inner_dim = key_i.shape[-1] + head_dim = inner_dim // self.query_num_heads + + query_i = query_i.view(batch_size, -1, self.query_num_heads, head_dim) + key_i = key_i.view(batch_size, -1, self.kv_num_heads, head_dim) + value_i = value_i.view(batch_size, -1, self.kv_num_heads, head_dim) + if hidden_states_masks is not None: + key_i = key_i * hidden_states_masks.view(batch_size, -1, 1, 1) + + if not self.single: + qkv_t, _ = self.to_qkv_t(encoder_hidden_states) + q_size_t = self.to_qkv_t.num_heads * self.head_dim + kv_size_t = self.to_qkv_t.num_kv_heads * self.head_dim + query_t, key_t, value_t = qkv_t.split([q_size_t, kv_size_t, kv_size_t], dim=-1) + query_t = self.q_rms_norm_t(query_t).to(dtype=dtype) + key_t = self.k_rms_norm_t(key_t).to(dtype=dtype) + value_t = value_t + + query_t = query_t.view(batch_size, -1, self.to_qkv_t.num_heads, head_dim) + key_t = key_t.view(batch_size, -1, self.to_qkv_t.num_kv_heads, head_dim) + value_t = value_t.view(batch_size, -1, self.to_qkv_t.num_kv_heads, head_dim) + + num_image_tokens = query_i.shape[1] + num_text_tokens = query_t.shape[1] + query = torch.cat([query_i, query_t], dim=1) + key = torch.cat([key_i, key_t], dim=1) + value = torch.cat([value_i, value_t], dim=1) + else: + query = query_i + key = key_i + value = value_i + + cos, sin = image_rotary_emb + cos = cos.squeeze(2).to(query.dtype) + sin = sin.squeeze(2).to(query.dtype) + if query.shape[-1] == image_rotary_emb[0].shape[-1] * 2: + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) + else: + query_1, query_2 = query.chunk(2, dim=-1) + key_1, key_2 = key.chunk(2, dim=-1) + key_1 = self.rope(query_1, cos, sin) + key_1 = self.rope(key_1, cos, sin) + query = torch.cat([query_1, query_2], dim=-1) + key = torch.cat([key_1, key_2], dim=-1) + + hidden_states = F.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.query_num_heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + if not self.single: + hidden_states_i, hidden_states_t = torch.split(hidden_states, [num_image_tokens, num_text_tokens], dim=1) + hidden_states_i = self.to_out(hidden_states_i) + hidden_states_t = self.to_out_t(hidden_states_t) + return hidden_states_i, hidden_states_t + else: + hidden_states = self.to_out(hidden_states) + return hidden_states + + +# Modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py +class MoEGate(nn.Module): + def __init__( + self, + embed_dim, + num_routed_experts=4, + num_activated_experts=2, + aux_loss_alpha=0.01, + _force_inference_output=False, + ): + super().__init__() + self.top_k = num_activated_experts + self.n_routed_experts = num_routed_experts + + self.scoring_func = "softmax" + self.alpha = aux_loss_alpha + self.seq_aux = False + + # topk selection algorithm + self.norm_topk_prob = False + self.gating_dim = embed_dim + self.weight = nn.Parameter(torch.randn(self.n_routed_experts, self.gating_dim) / embed_dim**0.5) + + self._force_inference_output = _force_inference_output + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + ### compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear(hidden_states, self.weight, None) + if self.scoring_func == "softmax": + scores = logits.softmax(dim=-1) + else: + raise NotImplementedError(f"insupportable scoring function for MoE gating: {self.scoring_func}") + + ### select top-k experts + topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False) + + ### norm gate to sum 1 + if self.top_k > 1 and self.norm_topk_prob: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + + ### expert-level computation auxiliary loss + if self.training and self.alpha > 0.0 and not self._force_inference_output: + scores_for_aux = scores + aux_topk = self.top_k + # always compute aux loss based on the naive greedy topk method + topk_idx_for_aux_loss = topk_idx.view(bsz, -1) + if self.seq_aux: + scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1) + ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device) + ce.scatter_add_( + 1, topk_idx_for_aux_loss, torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device) + ).div_(seq_len * aux_topk / self.n_routed_experts) + aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha + else: + mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts) + ce = mask_ce.float().mean(0) + + Pi = scores_for_aux.mean(0) + fi = ce * self.n_routed_experts + aux_loss = (Pi * fi).sum() * self.alpha + else: + aux_loss = None + return topk_idx, topk_weight, aux_loss + + +# Modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py +class MOEFeedForwardSwiGLU(nn.Module): + def __init__( + self, + dim: int, + hidden_dim: int, + num_routed_experts: int, + num_activated_experts: int, + _force_inference_output: bool = False, + ): + super().__init__() + self.shared_experts = HiDreamImageFeedForwardSwiGLU(dim, hidden_dim // 2) + self.experts = nn.ModuleList( + [HiDreamImageFeedForwardSwiGLU(dim, hidden_dim) for i in range(num_routed_experts)] + ) + self._force_inference_output = _force_inference_output + self.gate = MoEGate( + embed_dim=dim, + num_routed_experts=num_routed_experts, + num_activated_experts=num_activated_experts, + _force_inference_output=_force_inference_output, + ) + self.num_activated_experts = num_activated_experts + + def forward(self, x): + wtype = x.dtype + identity = x + orig_shape = x.shape + topk_idx, topk_weight, aux_loss = self.gate(x) + x = x.view(-1, x.shape[-1]) + flat_topk_idx = topk_idx.view(-1) + if self.training and not self._force_inference_output: + x = x.repeat_interleave(self.num_activated_experts, dim=0) + y = torch.empty_like(x, dtype=wtype) + for i, expert in enumerate(self.experts): + y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(dtype=wtype) + y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) + y = y.view(*orig_shape).to(dtype=wtype) + else: + y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape) + y = y + self.shared_experts(identity) + return y + + @torch.no_grad() + def moe_infer(self, x, flat_expert_indices, flat_expert_weights): + expert_cache = torch.zeros_like(x) + idxs = flat_expert_indices.argsort() + tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0) + token_idxs = idxs // self.num_activated_experts + + # TODO: This MoE expert routing is sequential and can be a bottleneck for models + # with MoE in every transformer block. Consider optimizing this in the future + # using techniques like grouped GEMM or batched expert execution. + for i, end_idx in enumerate(tokens_per_expert): + start_idx = 0 if i == 0 else tokens_per_expert[i - 1] + if start_idx == end_idx: + continue + expert = self.experts[i] + exp_token_idx = token_idxs[start_idx:end_idx] + expert_tokens = x[exp_token_idx] + expert_out = expert(expert_tokens) + expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]]) + + # for fp16 and other dtype + expert_cache = expert_cache.to(expert_out.dtype) + expert_cache.scatter_reduce_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out, reduce="sum") + return expert_cache + + +class TextProjection(nn.Module): + def __init__(self, in_features, hidden_size): + super().__init__() + self.linear = nn.Linear(in_features=in_features, out_features=hidden_size, bias=False) + + def forward(self, caption): + hidden_states = self.linear(caption) + return hidden_states + + +class HiDreamImageSingleTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + num_routed_experts: int = 4, + num_activated_experts: int = 2, + _force_inference_output: bool = False, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True)) + + # 1. Attention + self.norm1_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) + self.attn1 = HiDreamAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + single=True, + ) + + # 3. Feed-forward + self.norm3_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) + if num_routed_experts > 0: + self.ff_i = MOEFeedForwardSwiGLU( + dim=dim, + hidden_dim=4 * dim, + num_routed_experts=num_routed_experts, + num_activated_experts=num_activated_experts, + _force_inference_output=_force_inference_output, + ) + else: + self.ff_i = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim) + + def forward( + self, + hidden_states: torch.Tensor, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + wtype = hidden_states.dtype + shift_msa_i, scale_msa_i, gate_msa_i, shift_mlp_i, scale_mlp_i, gate_mlp_i = self.adaLN_modulation(temb)[ + :, None + ].chunk(6, dim=-1) + + # 1. MM-Attention + norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype) + norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i + attn_output_i = self.attn1( + norm_hidden_states, + hidden_states_masks, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = gate_msa_i * attn_output_i + hidden_states + + # 2. Feed-forward + norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i + ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states.to(dtype=wtype)) + hidden_states = ff_output_i + hidden_states + return hidden_states + + +class HiDreamImageTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + num_routed_experts: int = 4, + num_activated_experts: int = 2, + _force_inference_output: bool = False, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 12 * dim, bias=True)) + + # 1. Attention + self.norm1_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) + self.norm1_t = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) + self.attn1 = HiDreamAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + single=False, + ) + + # 3. Feed-forward + self.norm3_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) + if num_routed_experts > 0: + self.ff_i = MOEFeedForwardSwiGLU( + dim=dim, + hidden_dim=4 * dim, + num_routed_experts=num_routed_experts, + num_activated_experts=num_activated_experts, + _force_inference_output=_force_inference_output, + ) + else: + self.ff_i = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim) + self.norm3_t = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) + self.ff_t = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim) + + def forward( + self, + hidden_states: torch.Tensor, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + wtype = hidden_states.dtype + ( + shift_msa_i, + scale_msa_i, + gate_msa_i, + shift_mlp_i, + scale_mlp_i, + gate_mlp_i, + shift_msa_t, + scale_msa_t, + gate_msa_t, + shift_mlp_t, + scale_mlp_t, + gate_mlp_t, + ) = self.adaLN_modulation(temb)[:, None].chunk(12, dim=-1) + + # 1. MM-Attention + norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype) + norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i + norm_encoder_hidden_states = self.norm1_t(encoder_hidden_states).to(dtype=wtype) + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + scale_msa_t) + shift_msa_t + + attn_output_i, attn_output_t = self.attn1( + norm_hidden_states, + hidden_states_masks, + norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = gate_msa_i * attn_output_i + hidden_states + encoder_hidden_states = gate_msa_t * attn_output_t + encoder_hidden_states + + # 2. Feed-forward + norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i + norm_encoder_hidden_states = self.norm3_t(encoder_hidden_states).to(dtype=wtype) + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + scale_mlp_t) + shift_mlp_t + + ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states) + ff_output_t = gate_mlp_t * self.ff_t(norm_encoder_hidden_states) + hidden_states = ff_output_i + hidden_states + encoder_hidden_states = ff_output_t + encoder_hidden_states + return hidden_states, encoder_hidden_states + + +class HiDreamBlock(nn.Module): + def __init__(self, block: HiDreamImageTransformerBlock | HiDreamImageSingleTransformerBlock): + super().__init__() + self.block = block + + def forward( + self, + hidden_states: torch.Tensor, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self.block( + hidden_states=hidden_states, + hidden_states_masks=hidden_states_masks, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + +class HiDreamImageTransformer2DModel(nn.Module): + def __init__( + self, + od_config: OmniDiffusionConfig = None, + patch_size: int | None = None, + in_channels: int = 64, + out_channels: int | None = None, + num_layers: int = 16, + num_single_layers: int = 32, + attention_head_dim: int = 128, + num_attention_heads: int = 20, + caption_channels: list[int] = None, + text_emb_dim: int = 2048, + num_routed_experts: int = 4, + num_activated_experts: int = 2, + axes_dims_rope: tuple[int, int] = (32, 32), + max_resolution: tuple[int, int] = (128, 128), + llama_layers: list[int] = None, + force_inference_output: bool = False, + quant_config: "QuantizationConfig | None" = None, + ): + super().__init__() + self.parallel_config = od_config.parallel_config + + self.in_channels = in_channels + self.patch_size = patch_size + self.force_inference_output = force_inference_output + self.llama_layers = llama_layers + + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + self.t_embedder = HiDreamImageTimestepEmbed(self.inner_dim) + self.p_embedder = HiDreamImagePooledEmbed(text_emb_dim, self.inner_dim) + self.x_embedder = HiDreamImagePatchEmbed( + patch_size=patch_size, + in_channels=in_channels, + out_channels=self.inner_dim, + ) + self.pe_embedder = HiDreamImageEmbedND(theta=10000, axes_dim=axes_dims_rope) + + self.double_stream_blocks = nn.ModuleList( + [ + HiDreamBlock( + HiDreamImageTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_routed_experts=num_routed_experts, + num_activated_experts=num_activated_experts, + _force_inference_output=force_inference_output, + ) + ) + for _ in range(num_layers) + ] + ) + + self.single_stream_blocks = nn.ModuleList( + [ + HiDreamBlock( + HiDreamImageSingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_routed_experts=num_routed_experts, + num_activated_experts=num_activated_experts, + _force_inference_output=force_inference_output, + ) + ) + for _ in range(num_single_layers) + ] + ) + + self.final_layer = HiDreamImageOutEmbed(self.inner_dim, patch_size, self.out_channels) + + num_llama_projections = len(self.llama_layers) if self.llama_layers is not None else 0 + caption_channels = [caption_channels[1]] * num_llama_projections + [caption_channels[0]] + caption_projection = [] + for caption_channel in caption_channels: + caption_projection.append(TextProjection(in_features=caption_channel, hidden_size=self.inner_dim)) + self.caption_projection = nn.ModuleList(caption_projection) + self.max_seq = max_resolution[0] * max_resolution[1] // (patch_size * patch_size) + + def unpatchify(self, x: torch.Tensor, img_sizes: list[tuple[int, int]], is_training: bool) -> list[torch.Tensor]: + if is_training and not self.force_inference_output: + B, S, F = x.shape + C = F // (self.patch_size * self.patch_size) + x = ( + x.reshape(B, S, self.patch_size, self.patch_size, C) + .permute(0, 4, 1, 2, 3) + .reshape(B, C, S, self.patch_size * self.patch_size) + ) + else: + x_arr = [] + p1 = self.patch_size + p2 = self.patch_size + for i, img_size in enumerate(img_sizes): + pH, pW = img_size + t = x[i, : pH * pW].reshape(1, pH, pW, -1) + F_token = t.shape[-1] + C = F_token // (p1 * p2) + t = t.reshape(1, pH, pW, p1, p2, C) + t = t.permute(0, 5, 1, 3, 2, 4) + t = t.reshape(1, C, pH * p1, pW * p2) + x_arr.append(t) + x = torch.cat(x_arr, dim=0) + return x + + def patchify(self, hidden_states): + batch_size, channels, height, width = hidden_states.shape + patch_size = self.patch_size + patch_height, patch_width = height // patch_size, width // patch_size + device = hidden_states.device + dtype = hidden_states.dtype + + # create img_sizes + img_sizes = torch.tensor([patch_height, patch_width], dtype=torch.int64, device=device).reshape(-1) + img_sizes = img_sizes.unsqueeze(0).repeat(batch_size, 1) + + # create hidden_states_masks + if hidden_states.shape[-2] != hidden_states.shape[-1]: + hidden_states_masks = torch.zeros((batch_size, self.max_seq), dtype=dtype, device=device) + hidden_states_masks[:, : patch_height * patch_width] = 1.0 + else: + hidden_states_masks = None + + # create img_ids + img_ids = torch.zeros(patch_height, patch_width, 3, device=device) + row_indices = torch.arange(patch_height, device=device)[:, None] + col_indices = torch.arange(patch_width, device=device)[None, :] + img_ids[..., 1] = img_ids[..., 1] + row_indices + img_ids[..., 2] = img_ids[..., 2] + col_indices + img_ids = img_ids.reshape(patch_height * patch_width, -1) + + if hidden_states.shape[-2] != hidden_states.shape[-1]: + # Handle non-square latents + img_ids_pad = torch.zeros(self.max_seq, 3, device=device) + img_ids_pad[: patch_height * patch_width, :] = img_ids + img_ids = img_ids_pad.unsqueeze(0).repeat(batch_size, 1, 1) + else: + img_ids = img_ids.unsqueeze(0).repeat(batch_size, 1, 1) + + # patchify hidden_states + if hidden_states.shape[-2] != hidden_states.shape[-1]: + # Handle non-square latents + out = torch.zeros( + (batch_size, channels, self.max_seq, patch_size * patch_size), + dtype=dtype, + device=device, + ) + hidden_states = hidden_states.reshape( + batch_size, channels, patch_height, patch_size, patch_width, patch_size + ) + hidden_states = hidden_states.permute(0, 1, 2, 4, 3, 5) + hidden_states = hidden_states.reshape( + batch_size, channels, patch_height * patch_width, patch_size * patch_size + ) + out[:, :, 0 : patch_height * patch_width] = hidden_states + hidden_states = out + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch_size, self.max_seq, patch_size * patch_size * channels + ) + + else: + # Handle square latents + hidden_states = hidden_states.reshape( + batch_size, channels, patch_height, patch_size, patch_width, patch_size + ) + hidden_states = hidden_states.permute(0, 2, 4, 3, 5, 1) + hidden_states = hidden_states.reshape( + batch_size, patch_height * patch_width, patch_size * patch_size * channels + ) + + return hidden_states, hidden_states_masks, img_sizes, img_ids + + def forward( + self, + hidden_states: torch.Tensor, + timesteps: torch.LongTensor = None, + encoder_hidden_states_t5: torch.Tensor = None, + encoder_hidden_states_llama3: torch.Tensor = None, + pooled_embeds: torch.Tensor = None, + img_ids: torch.Tensor | None = None, + img_sizes: list[tuple[int, int]] | None = None, + hidden_states_masks: torch.Tensor | None = None, + return_dict: bool = True, + **kwargs, + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: + encoder_hidden_states = kwargs.get("encoder_hidden_states", None) + + if encoder_hidden_states is not None: + deprecation_message = ( + "The `encoder_hidden_states` argument is deprecated." + "Please use `encoder_hidden_states_t5` and `encoder_hidden_states_llama3` instead." + ) + deprecate("encoder_hidden_states", "0.35.0", deprecation_message) + encoder_hidden_states_t5 = encoder_hidden_states[0] + encoder_hidden_states_llama3 = encoder_hidden_states[1] + + if img_ids is not None and img_sizes is not None and hidden_states_masks is None: + deprecation_message = ( + "Passing `img_ids` and `img_sizes` with unpachified `hidden_states` is deprecated and will be ignored." + ) + deprecate("img_ids", "0.35.0", deprecation_message) + + if hidden_states_masks is not None and (img_ids is None or img_sizes is None): + raise ValueError("if `hidden_states_masks` is passed, `img_ids` and `img_sizes` must also be passed.") + elif hidden_states_masks is not None and hidden_states.ndim != 3: + raise ValueError( + "if `hidden_states_masks` is passed, `hidden_states` must be a 3D tensors with shape" + "(batch_size, patch_height * patch_width, patch_size * patch_size * channels)" + ) + + # spatial forward + batch_size = hidden_states.shape[0] + hidden_states_type = hidden_states.dtype + + # Patchify the input + if hidden_states_masks is None: + hidden_states, hidden_states_masks, img_sizes, img_ids = self.patchify(hidden_states) + + # Embed the hidden states + hidden_states = self.x_embedder(hidden_states) + + # 0. time + timesteps = self.t_embedder(timesteps, hidden_states_type) + p_embedder = self.p_embedder(pooled_embeds) + temb = timesteps + p_embedder + + encoder_hidden_states = [encoder_hidden_states_llama3[k] for k in self.llama_layers] + + if self.caption_projection is not None: + new_encoder_hidden_states = [] + for i, enc_hidden_state in enumerate(encoder_hidden_states): + enc_hidden_state = self.caption_projection[i](enc_hidden_state) + enc_hidden_state = enc_hidden_state.view(batch_size, -1, hidden_states.shape[-1]) + new_encoder_hidden_states.append(enc_hidden_state) + encoder_hidden_states = new_encoder_hidden_states + encoder_hidden_states_t5 = self.caption_projection[-1](encoder_hidden_states_t5) + encoder_hidden_states_t5 = encoder_hidden_states_t5.view(batch_size, -1, hidden_states.shape[-1]) + encoder_hidden_states.append(encoder_hidden_states_t5) + + txt_ids = torch.zeros( + batch_size, + encoder_hidden_states[-1].shape[1] + encoder_hidden_states[-2].shape[1] + encoder_hidden_states[0].shape[1], + 3, + device=img_ids.device, + dtype=img_ids.dtype, + ) + image_rotary_emb = self.pe_embedder(img_ids) + text_rotary_emb = self.pe_embedder(txt_ids) + concat_rotary_emb = ( + torch.cat([image_rotary_emb[..., 0, 0], text_rotary_emb[..., 0, 0]], dim=1), + torch.cat([image_rotary_emb[..., 1, 0], text_rotary_emb[..., 1, 0]], dim=1), + ) + + # 2. Blocks + block_id = 0 + initial_encoder_hidden_states = torch.cat([encoder_hidden_states[-1], encoder_hidden_states[-2]], dim=1) + initial_encoder_hidden_states_seq_len = initial_encoder_hidden_states.shape[1] + for bid, block in enumerate(self.double_stream_blocks): + cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id] + cur_encoder_hidden_states = torch.cat( + [initial_encoder_hidden_states, cur_llama31_encoder_hidden_states], dim=1 + ) + hidden_states, initial_encoder_hidden_states = block( + hidden_states=hidden_states, + hidden_states_masks=hidden_states_masks, + encoder_hidden_states=cur_encoder_hidden_states, + temb=temb, + image_rotary_emb=concat_rotary_emb, + ) + initial_encoder_hidden_states = initial_encoder_hidden_states[:, :initial_encoder_hidden_states_seq_len] + block_id += 1 + + image_tokens_seq_len = hidden_states.shape[1] + hidden_states = torch.cat([hidden_states, initial_encoder_hidden_states], dim=1) + hidden_states_seq_len = hidden_states.shape[1] + if hidden_states_masks is not None: + encoder_attention_mask_ones = torch.ones( + (batch_size, initial_encoder_hidden_states.shape[1] + cur_llama31_encoder_hidden_states.shape[1]), + device=hidden_states_masks.device, + dtype=hidden_states_masks.dtype, + ) + hidden_states_masks = torch.cat([hidden_states_masks, encoder_attention_mask_ones], dim=1) + + for bid, block in enumerate(self.single_stream_blocks): + cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id] + hidden_states = torch.cat([hidden_states, cur_llama31_encoder_hidden_states], dim=1) + + hidden_states = block( + hidden_states=hidden_states, + hidden_states_masks=hidden_states_masks, + encoder_hidden_states=None, + temb=temb, + image_rotary_emb=concat_rotary_emb, + ) + hidden_states = hidden_states[:, :hidden_states_seq_len] + block_id += 1 + + hidden_states = hidden_states[:, :image_tokens_seq_len, ...] + output = self.final_layer(hidden_states, temb) + output = self.unpatchify(output, img_sizes, self.training) + if hidden_states_masks is not None: + hidden_states_masks = hidden_states_masks[:, :image_tokens_seq_len] + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + # self-attn + (".to_qkv", ".to_q", "q"), + (".to_qkv", ".to_k", "k"), + (".to_qkv", ".to_v", "v"), + # self-attn no single + (".to_qkv_t", ".to_q_t", "q"), + (".to_qkv_t", ".to_k_t", "k"), + (".to_qkv_t", ".to_v_t", "v"), + # cross-attn + (".add_kv_proj", ".add_q_proj", "q"), + (".add_kv_proj", ".add_k_proj", "k"), + (".add_kv_proj", ".add_v_proj", "v"), + ] + + params_dict = dict(self.named_parameters()) + + # we need to load the buffers for beta and eps (XIELU) + for name, buffer in self.named_buffers(): + if name.endswith(".beta") or name.endswith(".eps"): + params_dict[name] = buffer + + loaded_params: set[str] = set() + for name, loaded_weight in weights: + original_name = name + lookup_name = name + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in original_name: + continue + lookup_name = original_name.replace(weight_name, param_name) + param = params_dict[lookup_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if lookup_name not in params_dict and ".to_out.0." in lookup_name: + lookup_name = lookup_name.replace(".to_out.0.", ".to_out.") + if lookup_name not in params_dict: + continue + param = params_dict[lookup_name] + if tp_size > 1 and any( + norm_name in lookup_name + for norm_name in [ + ".attn1.q_rms_norm", + ".attn1.k_rms_norm", + ] + ): + shard_size = loaded_weight.shape[0] // tp_size + loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(original_name) + loaded_params.add(lookup_name) + return loaded_params diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py new file mode 100644 index 000000000..39869a00d --- /dev/null +++ b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py @@ -0,0 +1,1027 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect +import json +import math +import os +from collections.abc import Callable, Iterable +from typing import Any + +import numpy as np +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.models import AutoencoderKL +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import deprecate, logging +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import ( + CLIPTextModelWithProjection, + CLIPTokenizer, + LlamaForCausalLM, + PreTrainedTokenizerFast, + T5EncoderModel, + T5Tokenizer, +) +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +logger = logging.get_logger(__name__) + + +def get_hidream_image_post_process_func( + od_config: OmniDiffusionConfig, +): + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + vae_config_path = os.path.join(model_path, "vae/config.json") + if not os.path.exists(vae_config_path): + raise FileNotFoundError( + f"VAE config not found at {vae_config_path}. " + "Please ensure the model path contains a valid VAE configuration." + ) + with open(vae_config_path) as f: + vae_config = json.load(f) + vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8 + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) + + def post_process_func( + images: torch.Tensor, + ): + return image_processor.postprocess(images) + + return post_process_func + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + timesteps (`torch.Tensor`): The timestep schedule from the scheduler. + num_inference_steps (`int`): The number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + dtype = getattr(od_config, "dtype", torch.bfloat16) + llama_path = self.od_config.extras["auxiliary_text_encoder"] + if llama_path is None: + logger.warning( + f"auxiliary_text_encoder is not provided. " + f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}" + ) + llama_path = self.od_config.extras["default_llama_model_id"] + logger.info(f"Loading LLAMA model from: {llama_path}") + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + ] + + self.device = get_local_device() + model = od_config.model + # Check if model is a local path + local_files_only = os.path.exists(model) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( + self.device + ) + self.text_encoder = CLIPTextModelWithProjection.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) + self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained( + model, subfolder="text_encoder_2", local_files_only=local_files_only + ) + self.tokenizer_2 = CLIPTokenizer.from_pretrained( + model, subfolder="tokenizer_2", local_files_only=local_files_only + ) + self.text_encoder_3 = T5EncoderModel.from_pretrained( + model, subfolder="text_encoder_3", local_files_only=local_files_only + ) + self.tokenizer_3 = T5Tokenizer.from_pretrained( + model, subfolder="tokenizer_3", local_files_only=local_files_only + ) + self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to( + self.device + ) + self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False) + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel) + self.transformer = HiDreamImageTransformer2DModel( + od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + ) + + self.stage = None + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + # HiDreamImage latents are turned into 2x2 patches and packed. + # This means the latent width and height has to be divisible + # by the patch size. So the vae scale factor is multiplied by the patch size to account for this + self.default_sample_size = 128 + if getattr(self, "tokenizer_4", None) is not None: + self.tokenizer_4.pad_token = self.tokenizer_4.eos_token + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def check_inputs( + self, + prompt, + prompt_2, + prompt_3, + prompt_4, + negative_prompt=None, + negative_prompt_2=None, + negative_prompt_3=None, + negative_prompt_4=None, + prompt_embeds_t5=None, + prompt_embeds_llama3=None, + negative_prompt_embeds_t5=None, + negative_prompt_embeds_llama3=None, + pooled_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + ): + if prompt is not None and pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. " + "Please make sure to only forward one of the two." + ) + elif prompt_2 is not None and pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}." + "Please make sure to only forward one of the two." + ) + elif prompt_3 is not None and prompt_embeds_t5 is not None: + raise ValueError( + f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}." + "Please make sure to only forward one of the two." + ) + elif prompt_4 is not None and prompt_embeds_llama3 is not None: + raise ValueError( + f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}." + "Please make sure to only forward one of the two." + ) + elif prompt is None and pooled_prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `pooled_prompt_embeds`." + "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined." + ) + elif prompt is None and prompt_embeds_t5 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_t5`." + "Cannot leave both `prompt` and `prompt_embeds_t5` undefined." + ) + elif prompt is None and prompt_embeds_llama3 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_llama3`." + "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): + raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): + raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") + elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)): + raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}") + + if negative_prompt is not None and negative_pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:" + f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:" + f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:" + f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two." + ) + elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:" + f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two." + ) + + if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None: + if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape: + raise ValueError( + "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`" + " must have the same shape when passed directly, but" + f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`" + f" {negative_pooled_prompt_embeds.shape}." + ) + if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None: + if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape: + raise ValueError( + "`prompt_embeds_t5` and `negative_prompt_embeds_t5`" + " must have the same shape when passed directly, but" + f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`" + f" {negative_prompt_embeds_t5.shape}." + ) + if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None: + if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape: + raise ValueError( + "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`" + " must have the same shape when passed directly, but" + f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`" + f" {negative_prompt_embeds_llama3.shape}." + ) + + def _get_t5_prompt_embeds( + self, + prompt: str | list[str] = None, + max_sequence_length: int = 128, + dtype: torch.dtype | None = None, + ): + dtype = dtype or self.text_encoder_3.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = self.tokenizer_3( + prompt, + padding="max_length", + max_length=min(max_sequence_length, self.tokenizer_3.model_max_length), + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_3.batch_decode( + untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1] + ) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}" + ) + + prompt_embeds = self.text_encoder_3( + text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device) + )[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) + return prompt_embeds + + def _get_clip_prompt_embeds( + self, + tokenizer, + text_encoder, + prompt: str | list[str], + max_sequence_length: int = 128, + dtype: torch.dtype | None = None, + ): + dtype = dtype or text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + model_max_length = self.tokenizer.model_max_length + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=min(max_sequence_length, model_max_length), + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {model_max_length} tokens: {removed_text}" + ) + prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True) + + # Use pooled output of CLIPTextModel + prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) + return prompt_embeds + + def _get_llama3_prompt_embeds( + self, + prompt: str | list[str] = None, + max_sequence_length: int = 128, + dtype: torch.dtype | None = None, + ): + dtype = dtype or self.text_encoder_4.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = self.tokenizer_4( + prompt, + padding="max_length", + max_length=min(max_sequence_length, self.tokenizer_4.model_max_length), + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_4.batch_decode( + untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1] + ) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}" + ) + + outputs = self.text_encoder_4( + text_input_ids.to(self.device), + attention_mask=attention_mask.to(self.device), + output_hidden_states=True, + ) + + prompt_embeds = outputs.hidden_states[1:] + prompt_embeds = torch.stack(prompt_embeds, dim=0) + return prompt_embeds + + def encode_prompt( + self, + prompt: str | list[str] | None = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + prompt_4: str | list[str] | None = None, + dtype: torch.dtype | None = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + negative_prompt_4: str | list[str] | None = None, + prompt_embeds_t5: list[torch.FloatTensor] | None = None, + prompt_embeds_llama3: list[torch.FloatTensor] | None = None, + negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None, + negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + max_sequence_length: int = 128, + lora_scale: float | None = None, + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = pooled_prompt_embeds.shape[0] + + if pooled_prompt_embeds is None: + pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( + self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype + ) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if len(negative_prompt) > 1 and len(negative_prompt) != batch_size: + raise ValueError(f"negative_prompt must be of length 1 or {batch_size}") + + negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( + self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype + ) + + if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1: + negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1) + + if pooled_prompt_embeds is None: + prompt_2 = prompt_2 or prompt + prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 + + if len(prompt_2) > 1 and len(prompt_2) != batch_size: + raise ValueError(f"prompt_2 must be of length 1 or {batch_size}") + + pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( + self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype + ) + + if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: + pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_prompt_2 = negative_prompt_2 or negative_prompt + negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 + + if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size: + raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}") + + negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( + self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype + ) + + if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: + negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1) + + if pooled_prompt_embeds is None: + pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_pooled_prompt_embeds = torch.cat( + [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1 + ) + + if prompt_embeds_t5 is None: + prompt_3 = prompt_3 or prompt + prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 + + if len(prompt_3) > 1 and len(prompt_3) != batch_size: + raise ValueError(f"prompt_3 must be of length 1 or {batch_size}") + + prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype) + + if prompt_embeds_t5.shape[0] == 1 and batch_size > 1: + prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) + + if do_classifier_free_guidance and negative_prompt_embeds_t5 is None: + negative_prompt_3 = negative_prompt_3 or negative_prompt + negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 + + if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size: + raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}") + + negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype) + + if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1: + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) + + if prompt_embeds_llama3 is None: + prompt_4 = prompt_4 or prompt + prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4 + + if len(prompt_4) > 1 and len(prompt_4) != batch_size: + raise ValueError(f"prompt_4 must be of length 1 or {batch_size}") + + prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype) + + if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + + if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None: + negative_prompt_4 = negative_prompt_4 or negative_prompt + negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4 + + if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size: + raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}") + + negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds( + negative_prompt_4, max_sequence_length, dtype + ) + + if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + + # duplicate pooled_prompt_embeds for each generation per prompt + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) + pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len, _ = prompt_embeds_t5.shape + if bs_embed == 1 and batch_size > 1: + prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}") + prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) + prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) + + # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt + _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape + if bs_embed == 1 and batch_size > 1: + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}") + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) + prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim) + + if do_classifier_free_guidance: + # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len = negative_pooled_prompt_embeds.shape + if bs_embed == 1 and batch_size > 1: + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}") + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt) + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape + if bs_embed == 1 and batch_size > 1: + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}") + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) + + # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt + _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape + if bs_embed == 1 and batch_size > 1: + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}") + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view( + -1, batch_size * num_images_per_prompt, seq_len, dim + ) + + return ( + prompt_embeds_t5, + negative_prompt_embeds_t5, + prompt_embeds_llama3, + negative_prompt_embeds_llama3, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) + + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + depr_message = ( + f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" + " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." + ) + deprecate( + "enable_vae_slicing", + "0.40.0", + depr_message, + ) + self.vae.enable_slicing() + + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + depr_message = ( + f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" + " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." + ) + deprecate( + "disable_vae_slicing", + "0.40.0", + depr_message, + ) + self.vae.disable_slicing() + + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + depr_message = ( + f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`" + " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." + ) + deprecate( + "enable_vae_tiling", + "0.40.0", + depr_message, + ) + self.vae.enable_tiling() + + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + depr_message = ( + f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is" + " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." + ) + deprecate( + "disable_vae_tiling", + "0.40.0", + depr_message, + ) + self.vae.disable_tiling() + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + generator, + latents=None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype) + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(self.device) + return latents + + def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len): + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + self.device, + sigmas=sigmas, + mu=mu, + ) + return timesteps, num_inference_steps + + def _extract_prompts(self, prompts): + """Extract prompt and negative_prompt from OmniPromptType list.""" + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts): + negative_prompt = None + elif prompts: + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts] + else: + negative_prompt = None + return prompt, negative_prompt + + def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool): + if get_classifier_free_guidance_world_size() == 1: + return True + + if true_cfg_scale <= 1: + logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.") + return False + + if not has_neg_prompt: + logger.warning( + "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings." + ) + return False + return True + + def diffuse( + self, + prompt_embeds_t5: torch.Tensor, + prompt_embeds_llama3: torch.Tensor, + pooled_prompt_embeds: torch.Tensor, + latents: torch.Tensor, + timesteps: torch.Tensor, + do_true_cfg: bool, + ) -> torch.Tensor: + with self.progress_bar(total=len(timesteps)) as pbar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + noise_pred = self.transformer( + hidden_states=latent_model_input, + timesteps=timestep, + encoder_hidden_states_t5=prompt_embeds_t5, + encoder_hidden_states_llama3=prompt_embeds_llama3, + pooled_embeds=pooled_prompt_embeds, + return_dict=False, + )[0] + noise_pred = -noise_pred + + # TODO: Modify CFG guidance + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() + + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + prompt_4: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + guidance_scale: float = 5.0, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + negative_prompt_4: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds_t5: torch.FloatTensor | None = None, + prompt_embeds_llama3: torch.FloatTensor | None = None, + negative_prompt_embeds_t5: torch.FloatTensor | None = None, + negative_prompt_embeds_llama3: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 128, + **kwargs, + ): + extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) + prompt = extracted_prompt or prompt + + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + sigmas = req.sampling_params.sigmas or sigmas + max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length + generator = req.sampling_params.generator or generator + true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale + if req.sampling_params.guidance_scale_provided: + guidance_scale = req.sampling_params.guidance_scale + num_images_per_prompt = ( + req.sampling_params.num_outputs_per_prompt + if req.sampling_params.num_outputs_per_prompt > 0 + else num_images_per_prompt + ) + + prompt_embeds = kwargs.get("prompt_embeds", None) + negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None) + + if prompt_embeds is not None: + deprecation_message = ( + "The `prompt_embeds` argument is deprecated." + " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead." + ) + deprecate("prompt_embeds", "0.35.0", deprecation_message) + prompt_embeds_t5 = prompt_embeds[0] + prompt_embeds_llama3 = prompt_embeds[1] + + if negative_prompt_embeds is not None: + deprecation_message = ( + "The `negative_prompt_embeds` argument is deprecated." + "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead." + ) + deprecate("negative_prompt_embeds", "0.35.0", deprecation_message) + negative_prompt_embeds_t5 = negative_prompt_embeds[0] + negative_prompt_embeds_llama3 = negative_prompt_embeds[1] + + division = self.vae_scale_factor * 2 + S_max = (self.default_sample_size * self.vae_scale_factor) ** 2 + scale = S_max / (width * height) + scale = math.sqrt(scale) + width, height = int(width * scale // division * division), int(height * scale // division * division) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + prompt_3, + prompt_4, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + negative_prompt_4=negative_prompt_4, + prompt_embeds_t5=prompt_embeds_t5, + prompt_embeds_llama3=prompt_embeds_llama3, + negative_prompt_embeds_t5=negative_prompt_embeds_t5, + negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + elif pooled_prompt_embeds is not None: + batch_size = pooled_prompt_embeds.shape[0] + + # TODO: CFG guidance configuration + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None + ) + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt + self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) + + # 3. Encode prompt + lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None + ( + prompt_embeds_t5, + negative_prompt_embeds_t5, + prompt_embeds_llama3, + negative_prompt_embeds_llama3, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + prompt_3=prompt_3, + prompt_4=prompt_4, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + negative_prompt_4=negative_prompt_4, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds_t5=prompt_embeds_t5, + prompt_embeds_llama3=prompt_embeds_llama3, + negative_prompt_embeds_t5=negative_prompt_embeds_t5, + negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + if self.do_classifier_free_guidance: + prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0) + prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1) + pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + pooled_prompt_embeds.dtype, + generator, + latents, + ) + + # 5. Prepare timesteps + timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) + self._num_timesteps = len(timesteps) + + # 6. Denoising loop + latents = self.diffuse( + prompt_embeds_t5, + prompt_embeds_llama3, + pooled_prompt_embeds, + latents, + timesteps, + do_true_cfg, + ) + + if output_type == "latent": + image = latents + + else: + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py new file mode 100644 index 000000000..0c5486ae5 --- /dev/null +++ b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py @@ -0,0 +1,1027 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect +import json +import math +import os +from collections.abc import Callable, Iterable +from typing import Any + +import numpy as np +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.models import AutoencoderKL +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import deprecate, logging +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import ( + CLIPTextModelWithProjection, + CLIPTokenizer, + LlamaForCausalLM, + PreTrainedTokenizerFast, + T5EncoderModel, + T5Tokenizer, +) +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +logger = logging.get_logger(__name__) + + +def get_hidream_image_post_process_func( + od_config: OmniDiffusionConfig, +): + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + vae_config_path = os.path.join(model_path, "vae/config.json") + if not os.path.exists(vae_config_path): + raise FileNotFoundError( + f"VAE config not found at {vae_config_path}. " + "Please ensure the model path contains a valid VAE configuration." + ) + with open(vae_config_path) as f: + vae_config = json.load(f) + vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8 + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) + + def post_process_func( + images: torch.Tensor, + ): + return image_processor.postprocess(images) + + return post_process_func + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + timesteps (`torch.Tensor`): The timestep schedule from the scheduler. + num_inference_steps (`int`): The number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class HiDreamImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + dtype = getattr(od_config, "dtype", torch.bfloat16) + llama_path = self.od_config.extras["auxiliary_text_encoder"] + if llama_path is None: + logger.warning( + f"auxiliary_text_encoder is not provided. " + f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}" + ) + llama_path = self.od_config.extras["default_llama_model_id"] + logger.info(f"Loading LLAMA model from: {llama_path}") + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + ] + + self.device = get_local_device() + model = od_config.model + # Check if model is a local path + local_files_only = os.path.exists(model) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( + self.device + ) + self.text_encoder = CLIPTextModelWithProjection.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) + self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained( + model, subfolder="text_encoder_2", local_files_only=local_files_only + ) + self.tokenizer_2 = CLIPTokenizer.from_pretrained( + model, subfolder="tokenizer_2", local_files_only=local_files_only + ) + self.text_encoder_3 = T5EncoderModel.from_pretrained( + model, subfolder="text_encoder_3", local_files_only=local_files_only + ) + self.tokenizer_3 = T5Tokenizer.from_pretrained( + model, subfolder="tokenizer_3", local_files_only=local_files_only + ) + self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to( + self.device + ) + self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False) + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel) + self.transformer = HiDreamImageTransformer2DModel( + od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + ) + + self.stage = None + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + # HiDreamImage latents are turned into 2x2 patches and packed. + # This means the latent width and height has to be divisible + # by the patch size. So the vae scale factor is multiplied by the patch size to account for this + self.default_sample_size = 128 + if getattr(self, "tokenizer_4", None) is not None: + self.tokenizer_4.pad_token = self.tokenizer_4.eos_token + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def check_inputs( + self, + prompt, + prompt_2, + prompt_3, + prompt_4, + negative_prompt=None, + negative_prompt_2=None, + negative_prompt_3=None, + negative_prompt_4=None, + prompt_embeds_t5=None, + prompt_embeds_llama3=None, + negative_prompt_embeds_t5=None, + negative_prompt_embeds_llama3=None, + pooled_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + ): + if prompt is not None and pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. " + "Please make sure to only forward one of the two." + ) + elif prompt_2 is not None and pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}." + "Please make sure to only forward one of the two." + ) + elif prompt_3 is not None and prompt_embeds_t5 is not None: + raise ValueError( + f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}." + "Please make sure to only forward one of the two." + ) + elif prompt_4 is not None and prompt_embeds_llama3 is not None: + raise ValueError( + f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}." + "Please make sure to only forward one of the two." + ) + elif prompt is None and pooled_prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `pooled_prompt_embeds`." + "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined." + ) + elif prompt is None and prompt_embeds_t5 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_t5`." + "Cannot leave both `prompt` and `prompt_embeds_t5` undefined." + ) + elif prompt is None and prompt_embeds_llama3 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_llama3`." + "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): + raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): + raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") + elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)): + raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}") + + if negative_prompt is not None and negative_pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:" + f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:" + f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:" + f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two." + ) + elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:" + f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two." + ) + + if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None: + if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape: + raise ValueError( + "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`" + " must have the same shape when passed directly, but" + f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`" + f" {negative_pooled_prompt_embeds.shape}." + ) + if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None: + if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape: + raise ValueError( + "`prompt_embeds_t5` and `negative_prompt_embeds_t5`" + " must have the same shape when passed directly, but" + f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`" + f" {negative_prompt_embeds_t5.shape}." + ) + if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None: + if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape: + raise ValueError( + "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`" + " must have the same shape when passed directly, but" + f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`" + f" {negative_prompt_embeds_llama3.shape}." + ) + + def _get_t5_prompt_embeds( + self, + prompt: str | list[str] = None, + max_sequence_length: int = 128, + dtype: torch.dtype | None = None, + ): + dtype = dtype or self.text_encoder_3.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = self.tokenizer_3( + prompt, + padding="max_length", + max_length=min(max_sequence_length, self.tokenizer_3.model_max_length), + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_3.batch_decode( + untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1] + ) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}" + ) + + prompt_embeds = self.text_encoder_3( + text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device) + )[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) + return prompt_embeds + + def _get_clip_prompt_embeds( + self, + tokenizer, + text_encoder, + prompt: str | list[str], + max_sequence_length: int = 128, + dtype: torch.dtype | None = None, + ): + dtype = dtype or text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + model_max_length = self.tokenizer.model_max_length + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=min(max_sequence_length, model_max_length), + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {model_max_length} tokens: {removed_text}" + ) + prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True) + + # Use pooled output of CLIPTextModel + prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) + return prompt_embeds + + def _get_llama3_prompt_embeds( + self, + prompt: str | list[str] = None, + max_sequence_length: int = 128, + dtype: torch.dtype | None = None, + ): + dtype = dtype or self.text_encoder_4.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = self.tokenizer_4( + prompt, + padding="max_length", + max_length=min(max_sequence_length, self.tokenizer_4.model_max_length), + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_4.batch_decode( + untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1] + ) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}" + ) + + outputs = self.text_encoder_4( + text_input_ids.to(self.device), + attention_mask=attention_mask.to(self.device), + output_hidden_states=True, + ) + + prompt_embeds = outputs.hidden_states[1:] + prompt_embeds = torch.stack(prompt_embeds, dim=0) + return prompt_embeds + + def encode_prompt( + self, + prompt: str | list[str] | None = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + prompt_4: str | list[str] | None = None, + dtype: torch.dtype | None = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + negative_prompt_4: str | list[str] | None = None, + prompt_embeds_t5: list[torch.FloatTensor] | None = None, + prompt_embeds_llama3: list[torch.FloatTensor] | None = None, + negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None, + negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + max_sequence_length: int = 128, + lora_scale: float | None = None, + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = pooled_prompt_embeds.shape[0] + + if pooled_prompt_embeds is None: + pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( + self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype + ) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if len(negative_prompt) > 1 and len(negative_prompt) != batch_size: + raise ValueError(f"negative_prompt must be of length 1 or {batch_size}") + + negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( + self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype + ) + + if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1: + negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1) + + if pooled_prompt_embeds is None: + prompt_2 = prompt_2 or prompt + prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 + + if len(prompt_2) > 1 and len(prompt_2) != batch_size: + raise ValueError(f"prompt_2 must be of length 1 or {batch_size}") + + pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( + self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype + ) + + if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: + pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_prompt_2 = negative_prompt_2 or negative_prompt + negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 + + if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size: + raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}") + + negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( + self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype + ) + + if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: + negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1) + + if pooled_prompt_embeds is None: + pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_pooled_prompt_embeds = torch.cat( + [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1 + ) + + if prompt_embeds_t5 is None: + prompt_3 = prompt_3 or prompt + prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 + + if len(prompt_3) > 1 and len(prompt_3) != batch_size: + raise ValueError(f"prompt_3 must be of length 1 or {batch_size}") + + prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype) + + if prompt_embeds_t5.shape[0] == 1 and batch_size > 1: + prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) + + if do_classifier_free_guidance and negative_prompt_embeds_t5 is None: + negative_prompt_3 = negative_prompt_3 or negative_prompt + negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 + + if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size: + raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}") + + negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype) + + if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1: + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) + + if prompt_embeds_llama3 is None: + prompt_4 = prompt_4 or prompt + prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4 + + if len(prompt_4) > 1 and len(prompt_4) != batch_size: + raise ValueError(f"prompt_4 must be of length 1 or {batch_size}") + + prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype) + + if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + + if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None: + negative_prompt_4 = negative_prompt_4 or negative_prompt + negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4 + + if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size: + raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}") + + negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds( + negative_prompt_4, max_sequence_length, dtype + ) + + if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + + # duplicate pooled_prompt_embeds for each generation per prompt + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) + pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len, _ = prompt_embeds_t5.shape + if bs_embed == 1 and batch_size > 1: + prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}") + prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) + prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) + + # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt + _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape + if bs_embed == 1 and batch_size > 1: + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}") + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) + prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim) + + if do_classifier_free_guidance: + # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len = negative_pooled_prompt_embeds.shape + if bs_embed == 1 and batch_size > 1: + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}") + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt) + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape + if bs_embed == 1 and batch_size > 1: + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}") + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) + + # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt + _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape + if bs_embed == 1 and batch_size > 1: + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}") + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view( + -1, batch_size * num_images_per_prompt, seq_len, dim + ) + + return ( + prompt_embeds_t5, + negative_prompt_embeds_t5, + prompt_embeds_llama3, + negative_prompt_embeds_llama3, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) + + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + depr_message = ( + f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" + " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." + ) + deprecate( + "enable_vae_slicing", + "0.40.0", + depr_message, + ) + self.vae.enable_slicing() + + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + depr_message = ( + f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" + " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." + ) + deprecate( + "disable_vae_slicing", + "0.40.0", + depr_message, + ) + self.vae.disable_slicing() + + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + depr_message = ( + f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`" + " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." + ) + deprecate( + "enable_vae_tiling", + "0.40.0", + depr_message, + ) + self.vae.enable_tiling() + + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + depr_message = ( + f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is" + " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." + ) + deprecate( + "disable_vae_tiling", + "0.40.0", + depr_message, + ) + self.vae.disable_tiling() + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + generator, + latents=None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype) + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(self.device) + return latents + + def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len): + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + self.device, + sigmas=sigmas, + mu=mu, + ) + return timesteps, num_inference_steps + + def _extract_prompts(self, prompts): + """Extract prompt and negative_prompt from OmniPromptType list.""" + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts): + negative_prompt = None + elif prompts: + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts] + else: + negative_prompt = None + return prompt, negative_prompt + + def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool): + if get_classifier_free_guidance_world_size() == 1: + return True + + if true_cfg_scale <= 1: + logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.") + return False + + if not has_neg_prompt: + logger.warning( + "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings." + ) + return False + return True + + def diffuse( + self, + prompt_embeds_t5: torch.Tensor, + prompt_embeds_llama3: torch.Tensor, + pooled_prompt_embeds: torch.Tensor, + latents: torch.Tensor, + timesteps: torch.Tensor, + do_true_cfg: bool, + ) -> torch.Tensor: + with self.progress_bar(total=len(timesteps)) as pbar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + noise_pred = self.transformer( + hidden_states=latent_model_input, + timesteps=timestep, + encoder_hidden_states_t5=prompt_embeds_t5, + encoder_hidden_states_llama3=prompt_embeds_llama3, + pooled_embeds=pooled_prompt_embeds, + return_dict=False, + )[0] + noise_pred = -noise_pred + + # TODO: Modify CFG guidance + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() + + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + prompt_4: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + guidance_scale: float = 5.0, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + negative_prompt_4: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds_t5: torch.FloatTensor | None = None, + prompt_embeds_llama3: torch.FloatTensor | None = None, + negative_prompt_embeds_t5: torch.FloatTensor | None = None, + negative_prompt_embeds_llama3: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 128, + **kwargs, + ): + extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) + prompt = extracted_prompt or prompt + + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + sigmas = req.sampling_params.sigmas or sigmas + max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length + generator = req.sampling_params.generator or generator + true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale + if req.sampling_params.guidance_scale_provided: + guidance_scale = req.sampling_params.guidance_scale + num_images_per_prompt = ( + req.sampling_params.num_outputs_per_prompt + if req.sampling_params.num_outputs_per_prompt > 0 + else num_images_per_prompt + ) + + prompt_embeds = kwargs.get("prompt_embeds", None) + negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None) + + if prompt_embeds is not None: + deprecation_message = ( + "The `prompt_embeds` argument is deprecated." + " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead." + ) + deprecate("prompt_embeds", "0.35.0", deprecation_message) + prompt_embeds_t5 = prompt_embeds[0] + prompt_embeds_llama3 = prompt_embeds[1] + + if negative_prompt_embeds is not None: + deprecation_message = ( + "The `negative_prompt_embeds` argument is deprecated." + "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead." + ) + deprecate("negative_prompt_embeds", "0.35.0", deprecation_message) + negative_prompt_embeds_t5 = negative_prompt_embeds[0] + negative_prompt_embeds_llama3 = negative_prompt_embeds[1] + + division = self.vae_scale_factor * 2 + S_max = (self.default_sample_size * self.vae_scale_factor) ** 2 + scale = S_max / (width * height) + scale = math.sqrt(scale) + width, height = int(width * scale // division * division), int(height * scale // division * division) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + prompt_3, + prompt_4, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + negative_prompt_4=negative_prompt_4, + prompt_embeds_t5=prompt_embeds_t5, + prompt_embeds_llama3=prompt_embeds_llama3, + negative_prompt_embeds_t5=negative_prompt_embeds_t5, + negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + elif pooled_prompt_embeds is not None: + batch_size = pooled_prompt_embeds.shape[0] + + # TODO: CFG guidance configuration + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None + ) + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt + self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) + + # 3. Encode prompt + lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None + ( + prompt_embeds_t5, + negative_prompt_embeds_t5, + prompt_embeds_llama3, + negative_prompt_embeds_llama3, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + prompt_3=prompt_3, + prompt_4=prompt_4, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + negative_prompt_4=negative_prompt_4, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds_t5=prompt_embeds_t5, + prompt_embeds_llama3=prompt_embeds_llama3, + negative_prompt_embeds_t5=negative_prompt_embeds_t5, + negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + if self.do_classifier_free_guidance: + prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0) + prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1) + pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + pooled_prompt_embeds.dtype, + generator, + latents, + ) + + # 5. Prepare timesteps + timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) + self._num_timesteps = len(timesteps) + + # 6. Denoising loop + latents = self.diffuse( + prompt_embeds_t5, + prompt_embeds_llama3, + pooled_prompt_embeds, + latents, + timesteps, + do_true_cfg, + ) + + if output_type == "latent": + image = latents + + else: + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index b6a5a3700..d8302c115 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -261,6 +261,11 @@ _DIFFUSION_MODELS = { "pipeline_diffusers_adapter", "DiffusersAdapterPipeline", ), + "HiDreamImagePipeline": ( + "hidream_image", + "pipeline_hidream_image", + "HiDreamImagePipeline", + ), } @@ -482,6 +487,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "OmniVoicePipeline": "get_omnivoice_post_process_func", "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", + "HiDreamImagePipeline": "get_hidream_image_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index ad8dbff6f..23ef9b855 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1950,6 +1950,10 @@ class AsyncOmniEngine: "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), + "extras": { + "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None), + "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), + }, **( { "profiler_config": asdict(kwargs["profiler_config"]) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index e4fddc830..cc6e9a4da 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -650,6 +650,16 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable AR stage profiler to include AR stage timing in stage_durations.", ) + + # Supplementary auxiliary text encoder parameters + # (e.g., the meta llama/meta llama-3.1-8b-instrument used by hidream) + omni_config_group.add_argument( + "--auxiliary-text-encoder", + type=str, + default=None, + help="Auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).", + ) + # Stash via type(self) so the docs hook (which execs this function in a # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError. type(self)._parser = serve_parser