424 lines
18 KiB
Python
424 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Diffusion pipeline for Ming-Image Design and Design-Layer checkpoints."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from collections.abc import Iterable
|
|
from typing import Any, ClassVar
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
|
from PIL import Image
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader
|
|
|
|
from vllm_omni.diffusion.compile import regionally_compile
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import (
|
|
DistributedAutoencoderKLQwenImage,
|
|
)
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.forward_context import (
|
|
set_forward_context_direct_condition,
|
|
set_forward_context_ref_latent,
|
|
)
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch
|
|
from vllm_omni.diffusion.models.ming_image.condition import MingImageConditioning
|
|
from vllm_omni.diffusion.models.ming_image.transformer import MingImageTransformer2DModel
|
|
from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
|
|
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DESIGN_PIPELINE = "MingImageDiffusionPipeline"
|
|
_LAYERED_PIPELINE = "MingImageLayeredDiffusionPipeline"
|
|
_VENDOR_TRANSFORMER_CLASS = "DiffusionTransformer"
|
|
_DIFFUSION_REQUIRED_PATTERNS = [
|
|
"model_index.json",
|
|
"scheduler/**",
|
|
"transformer/**",
|
|
"vae/**",
|
|
"mlp/**",
|
|
"connector/**",
|
|
]
|
|
|
|
|
|
def _validate_variant_config(
|
|
model_index: dict[str, Any] | None,
|
|
transformer_config: Any,
|
|
) -> bool:
|
|
"""Validate checkpoint-owned variant metadata and return layered mode."""
|
|
declared_transformer = getattr(transformer_config, "_class_name", None)
|
|
if declared_transformer != _VENDOR_TRANSFORMER_CLASS:
|
|
raise ValueError(
|
|
"Ming-Image transformer/config.json must preserve the vendor "
|
|
f"_class_name={_VENDOR_TRANSFORMER_CLASS!r}, got {declared_transformer!r}. "
|
|
"vLLM-Omni loads those weights through MingImageTransformer2DModel."
|
|
)
|
|
|
|
alignment_padding_mode = getattr(transformer_config, "alignment_padding_mode", None)
|
|
multi_frame_output = getattr(transformer_config, "multi_frame_output", None)
|
|
variant = (alignment_padding_mode, multi_frame_output)
|
|
if variant == ("zero_masked", False):
|
|
is_layer_decomposition = False
|
|
elif variant == ("learned", True):
|
|
is_layer_decomposition = True
|
|
else:
|
|
raise ValueError(
|
|
"Ming-Image transformer/config.json must use either "
|
|
"alignment_padding_mode='zero_masked' with multi_frame_output=False "
|
|
"or alignment_padding_mode='learned' with multi_frame_output=True; "
|
|
f"got alignment_padding_mode={alignment_padding_mode!r}, "
|
|
f"multi_frame_output={multi_frame_output!r}."
|
|
)
|
|
|
|
declared_pipeline = (model_index or {}).get("_class_name")
|
|
if declared_pipeline is not None:
|
|
if declared_pipeline not in {_DESIGN_PIPELINE, _LAYERED_PIPELINE}:
|
|
raise ValueError(
|
|
"Ming-Image model_index.json must declare "
|
|
f"{_DESIGN_PIPELINE!r} or {_LAYERED_PIPELINE!r}, got {declared_pipeline!r}."
|
|
)
|
|
expected_pipeline = _LAYERED_PIPELINE if is_layer_decomposition else _DESIGN_PIPELINE
|
|
if declared_pipeline != expected_pipeline:
|
|
raise ValueError(
|
|
f"Ming-Image model_index.json declares {declared_pipeline!r}, but "
|
|
f"transformer/config.json describes {expected_pipeline!r}."
|
|
)
|
|
return is_layer_decomposition
|
|
|
|
|
|
class MingImageDiffusionPipeline(ZImagePipeline):
|
|
"""Ming-Image component adapter around the canonical Z-Image loop."""
|
|
|
|
supports_request_batch = False
|
|
|
|
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
|
_encoder_modules: ClassVar[list[str]] = ["conditioning"]
|
|
_vae_modules: ClassVar[list[str]] = ["vae"]
|
|
|
|
@property
|
|
def do_classifier_free_guidance(self) -> bool:
|
|
# The vendor Ming-Image pipeline only enables CFG above 1.0
|
|
return self._guidance_scale > 1
|
|
|
|
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
|
|
del prefix
|
|
nn.Module.__init__(self)
|
|
|
|
model_path = od_config.model
|
|
if not os.path.exists(model_path):
|
|
model_path = download_weights_from_hf_specific(
|
|
model_name_or_path=model_path,
|
|
cache_dir=None,
|
|
allow_patterns=_DIFFUSION_REQUIRED_PATTERNS,
|
|
revision=od_config.revision,
|
|
require_all=True,
|
|
)
|
|
local_files_only = os.path.isdir(model_path)
|
|
dtype = od_config.dtype
|
|
|
|
self.od_config = od_config
|
|
self._execution_device = get_local_device()
|
|
self.device = self._execution_device
|
|
model_index = get_diffusion_model_index(model_path, revision=od_config.revision) or {}
|
|
transformer_config = od_config.tf_model_config
|
|
self.is_layer_decomposition = _validate_variant_config(
|
|
model_index,
|
|
transformer_config,
|
|
)
|
|
|
|
self.default_num_inference_steps = 12
|
|
self.default_guidance_scale = 2.0 if self.is_layer_decomposition else 1.0
|
|
self._num_frames_per_prompt = 1
|
|
self._pending_prompt_embeds: list[torch.Tensor] | None = None
|
|
self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None
|
|
self._uses_cudagraph_trees = False
|
|
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=model_path,
|
|
subfolder="transformer",
|
|
revision=od_config.revision,
|
|
prefix="transformer.",
|
|
fall_back_to_pt=True,
|
|
)
|
|
]
|
|
subfolders = ["scheduler", "transformer", "vae"]
|
|
|
|
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
|
model_path,
|
|
subfolder="scheduler",
|
|
local_files_only=local_files_only,
|
|
)
|
|
|
|
self.vae = from_pretrained_with_prefetch(
|
|
DistributedAutoencoderKLQwenImage.from_pretrained,
|
|
model_path,
|
|
subfolder="vae",
|
|
prefetch_list=subfolders,
|
|
local_files_only=local_files_only,
|
|
torch_dtype=dtype,
|
|
).to(self.device)
|
|
self.vae.eval()
|
|
vae_channels = int(getattr(self.vae.config, "input_channels", 3))
|
|
if vae_channels != 4:
|
|
raise ValueError(f"Ming-Image requires an RGBA VAE, got input_channels={vae_channels}.")
|
|
|
|
transformer_kwargs = get_transformer_config_kwargs(
|
|
od_config.tf_model_config,
|
|
MingImageTransformer2DModel,
|
|
)
|
|
self.transformer = MingImageTransformer2DModel(
|
|
quant_config=od_config.quantization_config,
|
|
**transformer_kwargs,
|
|
)
|
|
self.conditioning = MingImageConditioning(
|
|
model_path,
|
|
device=self.device,
|
|
dtype=dtype,
|
|
)
|
|
self.text_encoder = None
|
|
self.tokenizer = None
|
|
|
|
self.vae_scale_factor = 2 ** len(self.vae.config.temperal_downsample)
|
|
self.image_processor = VaeImageProcessor(
|
|
vae_scale_factor=self.vae_scale_factor * 2,
|
|
do_convert_rgb=False,
|
|
resample="bilinear",
|
|
)
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler
|
|
)
|
|
|
|
def setup_compile(self) -> None:
|
|
# Keep request preparation, scheduling, and VAE work eager,
|
|
# while capturing repeated DiT blocks for CUDAGraph Trees replay.
|
|
if self.od_config.diffusion_compile_granularity != "regional":
|
|
logger.warning(
|
|
"Ming-Image CUDA Graph uses regional DiT compilation; diffusion_compile_granularity=%r is ignored.",
|
|
self.od_config.diffusion_compile_granularity,
|
|
)
|
|
self.transformer = regionally_compile(
|
|
self.transformer,
|
|
mode="reduce-overhead",
|
|
fullgraph=True,
|
|
dynamic=self.od_config.diffusion_compile_dynamic,
|
|
)
|
|
self._uses_cudagraph_trees = True
|
|
|
|
def encode_prompt(self, *args, **kwargs):
|
|
del args, kwargs
|
|
if self._pending_prompt_embeds is None or self._pending_negative_prompt_embeds is None:
|
|
raise RuntimeError("Ming-Image conditioning is only available during forward.")
|
|
return self._pending_prompt_embeds, self._pending_negative_prompt_embeds
|
|
|
|
def prepare_latents(self, batch_size, *args, **kwargs):
|
|
frames = self._num_frames_per_prompt
|
|
flat = super().prepare_latents(batch_size * frames, *args, **kwargs)
|
|
return torch.stack(flat.chunk(frames, dim=0), dim=2)
|
|
|
|
def _encode_reference(
|
|
self,
|
|
reference: Any | None,
|
|
height: int,
|
|
width: int,
|
|
) -> torch.Tensor | None:
|
|
if reference is None:
|
|
return None
|
|
if isinstance(reference, list):
|
|
if len(reference) != 1:
|
|
raise ValueError("Ming-Image currently accepts exactly one reference image.")
|
|
reference = reference[0]
|
|
if isinstance(reference, Image.Image):
|
|
reference = reference.convert("RGBA")
|
|
if not isinstance(reference, torch.Tensor):
|
|
reference = self.image_processor.preprocess(reference, height=height, width=width)
|
|
if reference.ndim != 4 or reference.shape[1] != 4:
|
|
raise ValueError("Ming-Image reference must be a batched RGBA tensor or an RGBA-compatible image.")
|
|
reference = reference.to(device=self.device, dtype=self.vae.dtype).unsqueeze(2)
|
|
latent = self.vae.encode(reference).latent_dist.mode()
|
|
return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
|
|
|
|
@staticmethod
|
|
def _get_prompt_extra(req: DiffusionRequestBatch) -> dict[str, Any]:
|
|
if not req.prompts:
|
|
return {}
|
|
prompt = req.prompts[0]
|
|
if isinstance(prompt, dict):
|
|
return dict(prompt.get("extra") or {})
|
|
if hasattr(prompt, "_asdict"):
|
|
return dict(prompt._asdict().get("extra") or {})
|
|
return {}
|
|
|
|
def _configure_output_frames(
|
|
self,
|
|
*,
|
|
reference: Any | None,
|
|
num_layers: int,
|
|
is_dummy_run: bool,
|
|
) -> None:
|
|
if num_layers < 1:
|
|
raise ValueError("num_layers must be at least 1.")
|
|
if self.is_layer_decomposition:
|
|
if reference is None and not is_dummy_run:
|
|
raise ValueError("Ming-Image Design-Layer requires a reference image.")
|
|
self._num_frames_per_prompt = num_layers + 1
|
|
return
|
|
if num_layers != 1:
|
|
raise ValueError("Ming-Image Design supports exactly one output frame.")
|
|
self._num_frames_per_prompt = 1
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
|
|
sampling = req.sampling_params
|
|
if sampling.num_outputs_per_prompt != 1:
|
|
# TODO(yuanheng-zhao): enable after supporting batching
|
|
raise ValueError(
|
|
f"Ming-Image currently supports num_outputs_per_prompt=1 only, got {sampling.num_outputs_per_prompt}."
|
|
)
|
|
|
|
extra = self._get_prompt_extra(req)
|
|
extra_args = sampling.extra_args or {}
|
|
|
|
is_dummy_run = req.is_dummy_run()
|
|
query_hidden = extra.get("query_hidden_states")
|
|
direct_hidden = extra.get("direct_hidden_states")
|
|
if query_hidden is None or direct_hidden is None:
|
|
if not is_dummy_run:
|
|
raise ValueError("Ming-Image requests require query and direct conditions.")
|
|
logger.warning("Ming-Image conditions are absent during warmup; using zero tensors.")
|
|
query_hidden = torch.zeros((256, 2048), device=self.device, dtype=self.od_config.dtype)
|
|
direct_hidden = torch.zeros((1, 6144), device=self.device, dtype=self.od_config.dtype)
|
|
if not isinstance(query_hidden, torch.Tensor) or not isinstance(direct_hidden, torch.Tensor):
|
|
raise TypeError("Ming-Image query and direct conditions must be tensors.")
|
|
if query_hidden.ndim == 2:
|
|
query_hidden = query_hidden.unsqueeze(0)
|
|
if direct_hidden.ndim == 2:
|
|
direct_hidden = direct_hidden.unsqueeze(0)
|
|
query_hidden = query_hidden.to(device=self.device, dtype=self.od_config.dtype)
|
|
direct_hidden = direct_hidden.to(device=self.device, dtype=self.od_config.dtype)
|
|
cap_feats, direct_condition = self.conditioning(query_hidden, direct_hidden)
|
|
|
|
reference = extra.get("reference_image")
|
|
num_layers = int(extra_args.get("num_layers", extra.get("num_layers", 1)))
|
|
self._configure_output_frames(
|
|
reference=reference,
|
|
num_layers=num_layers,
|
|
is_dummy_run=is_dummy_run,
|
|
)
|
|
|
|
height = int(extra_args.get("height") or sampling.height or 1024)
|
|
width = int(extra_args.get("width") or sampling.width or 1024)
|
|
steps = int(sampling.num_inference_steps or self.default_num_inference_steps)
|
|
cfg = float(sampling.guidance_scale if sampling.guidance_scale is not None else self.default_guidance_scale)
|
|
seed = extra_args.get("seed", sampling.seed)
|
|
generator = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else sampling.generator
|
|
|
|
ref_latent = self._encode_reference(reference, height, width)
|
|
positive = [item for item in cap_feats]
|
|
negative = [torch.zeros_like(item) for item in positive]
|
|
self._pending_prompt_embeds = positive
|
|
self._pending_negative_prompt_embeds = negative
|
|
|
|
apply_cfg = cfg > 1
|
|
context_direct = (
|
|
torch.cat([direct_condition, torch.zeros_like(direct_condition)], dim=0) if apply_cfg else direct_condition
|
|
)
|
|
context_ref = ref_latent
|
|
if apply_cfg and context_ref is not None:
|
|
context_ref = context_ref.repeat(2, 1, 1, 1, 1)
|
|
|
|
inner_sampling = OmniDiffusionSamplingParams(
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=steps,
|
|
guidance_scale=cfg,
|
|
generator=generator,
|
|
output_type="latent",
|
|
)
|
|
inner_req = DiffusionRequestBatch(
|
|
requests=[
|
|
OmniDiffusionRequest(
|
|
prompt={"prompt": ""},
|
|
sampling_params=inner_sampling,
|
|
request_id=req.request_id or "ming-image",
|
|
)
|
|
]
|
|
)
|
|
|
|
set_forward_context_ref_latent(context_ref)
|
|
set_forward_context_direct_condition(context_direct)
|
|
try:
|
|
latent_output = super().forward(inner_req)
|
|
if not isinstance(latent_output.output, torch.Tensor):
|
|
raise TypeError("Ming-Image denoising must return latent tensors.")
|
|
image = self._decode_latent_frames(latent_output.output)
|
|
return DiffusionOutput(
|
|
output=image,
|
|
stage_durations=latent_output.stage_durations,
|
|
)
|
|
finally:
|
|
set_forward_context_ref_latent(None)
|
|
set_forward_context_direct_condition(None)
|
|
self._pending_prompt_embeds = None
|
|
self._pending_negative_prompt_embeds = None
|
|
self._num_frames_per_prompt = 1
|
|
|
|
@staticmethod
|
|
def _flatten_latent_frames(latents: torch.Tensor) -> torch.Tensor:
|
|
if latents.ndim == 4:
|
|
return latents
|
|
if latents.ndim != 5:
|
|
raise ValueError(f"Expected 4D or 5D Ming-Image latents, got {tuple(latents.shape)}")
|
|
batch, channels, frames, height, width = latents.shape
|
|
return latents.permute(2, 0, 1, 3, 4).reshape(
|
|
frames * batch,
|
|
channels,
|
|
height,
|
|
width,
|
|
)
|
|
|
|
def _decode_latent_frames(self, latents: torch.Tensor) -> torch.Tensor:
|
|
latents = self._flatten_latent_frames(latents).to(self.vae.dtype).unsqueeze(2)
|
|
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
|
|
image = self.vae.decode(latents, return_dict=False)[0]
|
|
if image.ndim == 5:
|
|
if image.shape[2] != 1:
|
|
raise ValueError(
|
|
f"Ming-Image VAE returned multiple decoded frames for a single latent frame: {tuple(image.shape)}"
|
|
)
|
|
image = image.squeeze(2)
|
|
return image
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
loaded = AutoWeightsLoader(self).load_weights(weights)
|
|
loaded |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
|
|
loaded |= {f"conditioning.{name}" for name, _ in self.conditioning.named_parameters()}
|
|
return loaded
|
|
|
|
|
|
def get_ming_image_post_process_func(od_config: OmniDiffusionConfig):
|
|
del od_config
|
|
image_processor = VaeImageProcessor(vae_scale_factor=16, do_convert_rgb=False)
|
|
|
|
def post_process(images: torch.Tensor):
|
|
if images.ndim == 5:
|
|
images = images.permute(0, 2, 1, 3, 4).flatten(0, 1)
|
|
return image_processor.postprocess(images.float())
|
|
|
|
return post_process
|
|
|
|
|
|
__all__ = ["MingImageDiffusionPipeline", "get_ming_image_post_process_func"]
|