diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 1b6334d3f..75e1028b3 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -28,6 +28,7 @@ steps: depends_on: image-build commands: - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py + - pytest -s -v tests/e2e/offline_inference/test_ovis_image.py agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 63ff9820e..3a739d1cd 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -24,6 +24,7 @@ th { | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | | `WanPipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | +| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | diff --git a/tests/e2e/offline_inference/test_ovis_image.py b/tests/e2e/offline_inference/test_ovis_image.py new file mode 100644 index 000000000..9238ab039 --- /dev/null +++ b/tests/e2e/offline_inference/test_ovis_image.py @@ -0,0 +1,278 @@ +""" +Tests for Ovis Image model pipeline. + +Strategy: +1. `mock_dependencies` fixture mocks heavy external components (VAE, Scheduler, TextEncoder) + to allow fast testing of the pipeline logic without downloading weights. + - Mocks are configured to return tensors on the correct device. + - Transformer is mocked dynamically to return random noise of correct shape. + +2. `test_real_transformer_init_and_forward` tests the actual `OvisImageTransformer2DModel` + initialization and forward pass with a small configuration to ensure code coverage + and correctness of the model definition itself, independent of the pipeline mocks. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig + +# Mock the OvisImageTransformer2DModel to avoid complex init if needed, +# or let it run if it's lightweight. It's likely not lightweight. +# Better to mock the transformer forwarding to return random noise. +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image import OvisImagePipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest + + +@pytest.fixture +def mock_dependencies(monkeypatch): + """ + Mock external dependencies to avoid loading real models. + """ + device = get_local_device() + + # Mock Tokenizer + mock_tokenizer = MagicMock() + mock_tokenizer.return_value = MagicMock( + input_ids=torch.zeros((1, 50), dtype=torch.long, device=device), + attention_mask=torch.ones((1, 50), dtype=torch.long, device=device), + ) + mock_tokenizer.apply_chat_template.return_value = "dummy prompt" + mock_tokenizer.model_max_length = 1024 + + # Mock Text Encoder + mock_text_encoder = MagicMock() + mock_text_encoder.dtype = torch.float32 + # Output of text encoder must be on the same device as inputs (which are moved to execution_device) + mock_text_encoder.return_value.last_hidden_state = torch.randn(1, 50, 32, device=device) + + # Mock VAE + mock_vae = MagicMock() + mock_vae.config.block_out_channels = [128, 256, 512, 512] # Scale factor 8 + mock_vae.config.scale_factor_temporal = 1 + mock_vae.config.scale_factor_spatial = 8 + mock_vae.config.scaling_factor = 0.18215 + mock_vae.config.shift_factor = 0.0 + # Decode return value + mock_vae.decode.return_value = [torch.randn(1, 3, 128, 128, device=device)] + # Ensure .to() returns self so configuration persists + mock_vae.to.return_value = mock_vae + + # Mock Scheduler + mock_scheduler = MagicMock() + mock_scheduler.config = MagicMock() + # Timesteps on device to match latents during denoising loop interaction if needed + mock_scheduler.timesteps = torch.tensor([1.0, 0.5, 0.0], device=device) + mock_scheduler.set_timesteps.return_value = None + + # Make step return dynamic based on input sample shape + def mock_scheduler_step(model_output, timestep, sample, **kwargs): + # sample is the latents, should be preserved + return (torch.randn_like(sample),) + + mock_scheduler.step.side_effect = mock_scheduler_step + + module_path = "vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image" + + monkeypatch.setattr(f"{module_path}.Qwen2TokenizerFast.from_pretrained", lambda *a, **k: mock_tokenizer) + monkeypatch.setattr(f"{module_path}.Qwen3Model.from_pretrained", lambda *a, **k: mock_text_encoder) + monkeypatch.setattr(f"{module_path}.AutoencoderKL.from_pretrained", lambda *a, **k: mock_vae) + monkeypatch.setattr( + f"{module_path}.FlowMatchEulerDiscreteScheduler.from_pretrained", lambda *a, **k: mock_scheduler + ) + + return { + "tokenizer": mock_tokenizer, + "text_encoder": mock_text_encoder, + "vae": mock_vae, + "scheduler": mock_scheduler, + "device": device, + } + + +@pytest.fixture +def ovis_pipeline(mock_dependencies, monkeypatch): + """ + Creates an OvisImagePipeline instance with mocked components. + """ + # Create config + tf_config = TransformerConfig( + params={ + "in_channels": 4, + "out_channels": 4, + "sample_size": 32, + "patch_size": 2, + "num_attention_heads": 4, + "attention_head_dim": 8, + "num_layers": 1, + "caption_channels": 32, + } + ) + + od_config = OmniDiffusionConfig( + model="dummy-ovis", + tf_model_config=tf_config, + dtype=torch.float32, + num_gpus=1, + ) + + # Mock Transformer Layer separately to avoid full init + # We patch OvisImageTransformer2DModel class in the module + mock_transformer_cls = MagicMock() + mock_transformer_instance = MagicMock() + mock_transformer_instance.dtype = torch.float32 + mock_transformer_instance.in_channels = 16 # Must be 16 so num_channel_latents=4, packed=16 + # Forward return: noise prediction + + def mock_forward(hidden_states, *args, **kwargs): + # hidden_states shape: (B, SeqLen, Channels) + return (torch.randn_like(hidden_states),) + + mock_transformer_instance.forward.side_effect = mock_forward + # Also make the instance itself callable to mimic __call__ + mock_transformer_instance.side_effect = mock_forward + + mock_transformer_cls.return_value = mock_transformer_instance + + monkeypatch.setattr( + "vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image.OvisImageTransformer2DModel", mock_transformer_cls + ) + + # Initialize pipeline + # We use a dummy model path check override + with patch("os.path.exists", return_value=True): + pipeline = OvisImagePipeline(od_config=od_config) + + return pipeline + + +def test_interface_compliance(ovis_pipeline): + """Verify methods required by vllm-omni framework.""" + assert hasattr(ovis_pipeline, "load_weights") + assert hasattr(ovis_pipeline, "scheduler") + assert hasattr(ovis_pipeline, "transformer") + assert hasattr(ovis_pipeline, "text_encoder") + # assert hasattr(ovis_pipeline, "vae") # Ovis uses VAE + + +def test_basic_generation(ovis_pipeline): + """Test the forward pass logic.""" + # Setup request + req = OmniDiffusionRequest( + prompt="A photo of a cat", + height=256, + width=256, + num_inference_steps=2, + guidance_scale=1.0, + ) + + output = ovis_pipeline(req) + + assert output is not None + assert output.output is not None + # Output should be a tensor from mocked VAE decode [torch.randn(1, 3, 128, 128)] + assert isinstance(output.output, torch.Tensor) + assert output.output.shape == (1, 3, 128, 128) + + # Check that transformer was called + assert ovis_pipeline.transformer.call_count > 0 + + +def test_guidance_scale(ovis_pipeline): + """Test that classifier-free guidance path is taken when scale > 1.0.""" + req = OmniDiffusionRequest( + prompt="A photo of a cat", + negative_prompt="bad quality", + height=256, + width=256, + num_inference_steps=1, + guidance_scale=2.0, # Trigger CFG + ) + + ovis_pipeline(req) + assert ovis_pipeline.transformer.call_count >= 2 + + +def test_resolution_check(ovis_pipeline): + """Test resolution divisible validation logic if present.""" + # Pass odd resolution + req = OmniDiffusionRequest( + prompt="test", + height=250, # Not divisible by 16 (8*2) + width=250, + ) + + # Should warn but proceed (as per code I read earlier) or resize? + # The code had `logger.warning(...)` + + output = ovis_pipeline(req) + assert output is not None + + +def test_real_transformer_init_and_forward(): + """Test the real OvisImageTransformer2DModel initialization and forward pass for coverage.""" + from unittest.mock import patch + + from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel + + device = get_local_device() + tf_config = TransformerConfig( + params={ + "patch_size": 2, + "in_channels": 16, + "out_channels": 16, + "num_layers": 1, + "num_single_layers": 1, + "attention_head_dim": 8, + "num_attention_heads": 2, + "joint_attention_dim": 32, + "axes_dims_rope": (4, 4, 4), + } + ) + + od_config = OmniDiffusionConfig(model="dummy-ovis", tf_model_config=tf_config, dtype=torch.float32, num_gpus=1) + + # Mock distributed state for QKVParallelLinear initialization + # We patch get_tp_group because get_tensor_model_parallel_rank calls it and asserts _TP is not None + mock_group = MagicMock() + mock_group.rank_in_group = 0 + mock_group.world_size = 1 + + with patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group): + # Initialize real model + model = OvisImageTransformer2DModel( + od_config=od_config, + patch_size=1, + in_channels=16, + out_channels=16, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=32, + axes_dims_rope=(2, 2, 4), + ).to(device) + + # Create dummy inputs + B, Seq, C = 1, 16, 16 + hidden_states = torch.randn(B, Seq, C, device=device) + encoder_hidden_states = torch.randn(B, 10, 32, device=device) # joint_attention_dim=32 + timestep = torch.tensor([1], device=device) + img_ids = torch.zeros(Seq, 3, device=device) + txt_ids = torch.zeros(10, 3, device=device) + + # Run forward + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + return_dict=False, + ) + + assert output is not None + assert isinstance(output, tuple) + assert output[0].shape == hidden_states.shape diff --git a/vllm_omni/diffusion/models/ovis_image/__init__.py b/vllm_omni/diffusion/models/ovis_image/__init__.py new file mode 100644 index 000000000..1f7beae57 --- /dev/null +++ b/vllm_omni/diffusion/models/ovis_image/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ovis Image 7B diffusion model components.""" + +from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import ( + OvisImageTransformer2DModel, +) +from vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image import ( + OvisImagePipeline, + get_ovis_image_post_process_func, +) + +__all__ = [ + "OvisImagePipeline", + "OvisImageTransformer2DModel", + "get_ovis_image_post_process_func", +] diff --git a/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py b/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py new file mode 100644 index 000000000..ea3916b7a --- /dev/null +++ b/vllm_omni/diffusion/models/ovis_image/ovis_image_transformer.py @@ -0,0 +1,534 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2025 Alibaba Ovis-Image Team and The HuggingFace. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from diffusers.models.attention import FeedForward +from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle +from diffusers.utils import is_torch_npu_available +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + + +class OvisImageAttention(nn.Module): + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, + out_bias: bool = True, + eps: float = 1e-5, + out_dim: int = None, + context_pre_only: bool | None = None, + pre_only: bool = False, + ): + 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.use_bias = bias + self.dropout = dropout + self.out_dim = out_dim if out_dim is not None else query_dim + self.context_pre_only = context_pre_only + self.pre_only = pre_only + self.heads = out_dim // dim_head if out_dim is not None else heads + self.added_kv_proj_dim = added_kv_proj_dim + self.added_proj_bias = added_proj_bias + + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + + self.to_qkv = QKVParallelLinear( + hidden_size=query_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + disable_tp=True, + bias=bias, + ) + + if not self.pre_only: + self.to_out = nn.ModuleList([]) + self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(nn.Dropout(dropout)) + + if self.added_kv_proj_dim is not None: + self.norm_added_q = RMSNorm(dim_head, eps=eps) + self.norm_added_k = RMSNorm(dim_head, eps=eps) + + self.add_kv_proj = QKVParallelLinear( + hidden_size=self.added_kv_proj_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + disable_tp=True, + bias=added_proj_bias, + ) + + self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias) + + self.attn = Attention( + num_heads=heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + qkv, _ = self.to_qkv(hidden_states) + + query, key, value = qkv.chunk(3, dim=-1) + + query = query.unflatten(-1, (self.heads, -1)) + key = key.unflatten(-1, (self.heads, -1)) + value = value.unflatten(-1, (self.heads, -1)) + + query = self.norm_q(query) + key = self.norm_k(key) + + if self.added_kv_proj_dim is not None: + encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states) + encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1) + + encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) + encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) + encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) + + encoder_query = self.norm_added_q(encoder_query) + encoder_key = self.norm_added_k(encoder_key) + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + hidden_states = self.attn( + query, + key, + value, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + if encoder_hidden_states is not None: + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 + ) + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states) + + return hidden_states, encoder_hidden_states + else: + return hidden_states + + +class OvisImageSingleTransformerBlock(nn.Module): + def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0): + super().__init__() + self.mlp_hidden_dim = int(dim * mlp_ratio) + + self.norm = AdaLayerNormZeroSingle(dim) + self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim * 2) + self.act_mlp = nn.SiLU() + self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim) + + self.attn = OvisImageAttention( + query_dim=dim, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + bias=True, + eps=1e-6, + pre_only=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + residual = hidden_states + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + mlp_hidden_states, mlp_hidden_gate = torch.split( + self.proj_mlp(norm_hidden_states), [self.mlp_hidden_dim, self.mlp_hidden_dim], dim=-1 + ) + mlp_hidden_states = self.act_mlp(mlp_hidden_gate) * mlp_hidden_states + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + gate = gate.unsqueeze(1) + hidden_states = gate * self.proj_out(hidden_states) + hidden_states = residual + hidden_states + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:] + return encoder_hidden_states, hidden_states + + +class OvisImageTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + qk_norm: str = "rms_norm", + eps: float = 1e-6, + ): + super().__init__() + + self.norm1 = AdaLayerNormZero(dim) + self.norm1_context = AdaLayerNormZero(dim) + + self.attn = OvisImageAttention( + query_dim=dim, + added_kv_proj_dim=dim, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + context_pre_only=False, + bias=True, + eps=eps, + ) + + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="swiglu") + + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="swiglu") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( + encoder_hidden_states, emb=temb + ) + joint_attention_kwargs = joint_attention_kwargs or {} + + # Attention. + attention_outputs = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + if len(attention_outputs) == 2: + attn_output, context_attn_output = attention_outputs + elif len(attention_outputs) == 3: + attn_output, context_attn_output, ip_attn_output = attention_outputs + + # Process attention outputs for the `hidden_states`. + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = hidden_states + attn_output + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + ff_output = self.ff(norm_hidden_states) + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = hidden_states + ff_output + if len(attention_outputs) == 3: + hidden_states = hidden_states + ip_attn_output + + # Process attention outputs for the `encoder_hidden_states`. + context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output + encoder_hidden_states = encoder_hidden_states + context_attn_output + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +class OvisImagePosEmbed(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] + cos_out = [] + sin_out = [] + pos = ids.float() + is_mps = ids.device.type == "mps" + is_npu = ids.device.type == "npu" + freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 + for i in range(n_axes): + cos, sin = get_1d_rotary_pos_embed( + self.axes_dim[i], + pos[:, i], + theta=self.theta, + repeat_interleave_real=True, + use_real=True, + freqs_dtype=freqs_dtype, + ) + cos_out.append(cos) + sin_out.append(sin) + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +class OvisImageTransformer2DModel(nn.Module): + """ + The Transformer model introduced in Ovis-Image. + + Reference: https://github.com/AIDC-AI/Ovis-Image + + Args: + patch_size (`int`, defaults to `1`): + Patch size to turn the input data into small patches. + in_channels (`int`, defaults to `64`): + The number of channels in the input. + out_channels (`int`, *optional*, defaults to `None`): + The number of channels in the output. If not specified, it defaults to `in_channels`. + num_layers (`int`, defaults to `6`): + The number of layers of dual stream DiT blocks to use. + num_single_layers (`int`, defaults to `27`): + The number of layers of single stream DiT blocks to use. + attention_head_dim (`int`, defaults to `128`): + The number of dimensions to use for each attention head. + num_attention_heads (`int`, defaults to `24`): + The number of attention heads to use. + joint_attention_dim (`int`, defaults to `2048`): + The number of dimensions to use for the joint attention (embedding/channel dimension of + `encoder_hidden_states`). + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): + The dimensions to use for the rotary positional embeddings. + """ + + def __init__( + self, + od_config: OmniDiffusionConfig, + patch_size: int = 1, + in_channels: int = 64, + out_channels: int | None = 64, + num_layers: int = 6, + num_single_layers: int = 27, + attention_head_dim: int = 128, + num_attention_heads: int = 24, + joint_attention_dim: int = 2048, + axes_dims_rope: tuple[int] = (16, 56, 56), + ): + super().__init__() + model_config = od_config.tf_model_config + num_layers = model_config.num_layers + self.in_channels = in_channels + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + self.pos_embed = OvisImagePosEmbed(theta=10000, axes_dim=axes_dims_rope) + + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim) + + self.context_embedder_norm = RMSNorm(joint_attention_dim, eps=1e-6) + self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) + self.x_embedder = nn.Linear(in_channels, self.inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + OvisImageTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + ) + for _ in range(num_layers) + ] + ) + + self.single_transformer_blocks = nn.ModuleList( + [ + OvisImageSingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + ) + for _ in range(num_single_layers) + ] + ) + self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + timestep: torch.LongTensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + """ + The [`OvisImageTransformer2DModel`] forward method. + + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): + Input `hidden_states`. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): + Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. + timestep (`torch.LongTensor`): + Used to indicate denoising step. + img_ids: (`torch.Tensor`): + The position ids for image tokens. + txt_ids (`torch.Tensor`): + The position ids for text tokens. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + + hidden_states = self.x_embedder(hidden_states) + timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000 + + timesteps_proj = self.time_proj(timestep) + temb = self.timestep_embedder(timesteps_proj.to(device=hidden_states.device, dtype=hidden_states.dtype)) + + encoder_hidden_states = self.context_embedder_norm(encoder_hidden_states) + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + if txt_ids.ndim == 3: + logger.warning( + "Passing `txt_ids` 3d torch.Tensor is deprecated." + "Please remove the batch dimension and pass it as a 2d torch Tensor" + ) + txt_ids = txt_ids[0] + if img_ids.ndim == 3: + logger.warning( + "Passing `img_ids` 3d torch.Tensor is deprecated." + "Please remove the batch dimension and pass it as a 2d torch Tensor" + ) + img_ids = img_ids[0] + + ids = torch.cat((txt_ids, img_ids), dim=0) + if is_torch_npu_available(): + freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) + image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) + else: + image_rotary_emb = self.pos_embed(ids) + + for index_block, block in enumerate(self.transformer_blocks): + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + for index_block, block in enumerate(self.single_transformer_blocks): + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # self attn + (".to_qkv", ".to_q", "q"), + (".to_qkv", ".to_k", "k"), + (".to_qkv", ".to_v", "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: + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py new file mode 100644 index 000000000..b115afc93 --- /dev/null +++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py @@ -0,0 +1,728 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2025 Alibaba Ovis-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import json +import os +from collections.abc import Callable, Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from diffusers.image_processor import VaeImageProcessor +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from diffusers.utils.torch_utils import randn_tensor +from transformers import Qwen2TokenizerFast, Qwen3Model +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +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.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +logger = init_logger(__name__) + + +def get_ovis_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") + 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 + + +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 + + +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, +) -> tuple[torch.Tensor, int]: + r""" + Calls the scheduler's `set_timesteps` method and retrieves timetemps + from the scheduler after the call. Handles custom timeteps. Any kwargs will be supplied to `scheduler.set_timeteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`, *optional*): + 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: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is 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: + accepts_timesteps = "sigmas" 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" sigma 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 OvisImagePipeline( + nn.Module, +): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + self._execution_device = get_local_device() + model = od_config.model + local_files_only = os.path.exists(model) + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + + self.text_encoder = Qwen3Model.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) + + self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( + self._execution_device + ) + + self.tokenizer = Qwen2TokenizerFast.from_pretrained( + model, subfolder="tokenizer", local_files_only=local_files_only + ) + + self.transformer = OvisImageTransformer2DModel(od_config=od_config) + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + + self.tokenizer_max_length = 1024 + self.system_prompt = """Describe the image by detailing the color, quantity, text, shape, size, texture, spatial + relationships of the objects and background: """ + self.user_prompt_begin_id = 28 + self.tokenizer_max_length = 256 + self.user_prompt_begin_id + self.default_sample_size = 128 + + def _get_messages( + self, + prompt: str | list[str] = None, + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + + messages = [] + + for each_prompt in prompt: + message = [ + { + "role": "user", + "content": self.system_prompt + each_prompt, + } + ] + message = self.tokenizer.apply_chat_template( + message, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + messages.append(message) + + return messages + + def _get_ovis_prompt_embeds( + self, + prompt: str | list[str] = None, + num_images_per_prompt: int = 1, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + messages = self._get_messages(prompt) + + batch_size = len(messages) + + tokens = self.tokenizer( + messages, + padding="max_length", + truncation=True, + max_length=self.tokenizer_max_length, + return_tensors="pt", + add_special_tokens=False, + ) + + input_ids = tokens.input_ids.to(device=device) + attention_mask = tokens.attention_mask.to(device=device) + + outputs = self.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + ) + + prompt_embeds = outputs.last_hidden_state + prompt_embeds = prompt_embeds * attention_mask[..., None] + prompt_embeds = prompt_embeds[:, self.user_prompt_begin_id :, :] + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + return prompt_embeds + + def encode_prompt( + self, + prompt: str | list[str], + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.FloatTensor | None = None, + ): + r""" + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + device: (`torch.device`, *optional*): + torch.device + num_images_per_prompt: (`int`): + number of images that should be generated per prompt + prompt_embeds: (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, text embeddings will be generated from `prompt` input argument. + """ + + device = device or self._execution_device + + if prompt_embeds is None: + prompt_embeds = self._get_ovis_prompt_embeds( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + ) + + dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype + text_ids = torch.zeros(prompt_embeds.shape[1], 3) + text_ids[..., 1] = text_ids[..., 1] + torch.arange(prompt_embeds.shape[1])[None, :] + text_ids[..., 2] = text_ids[..., 2] + torch.arange(prompt_embeds.shape[1])[None, :] + text_ids = text_ids.to(device=device, dtype=dtype) + return prompt_embeds, text_ids + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + max_sequence_length=None, + ): + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"""`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are + {height} and {width}. Dimension will be resized accordingly""" + ) + + # if callback_on_step_end_tensor_inputs is not None and not all( + # k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + # ): + # raise ValueError( + # f"""`callback_on_step_end_tensor_inputs` has to contain the following keys: + # {self._callback_tensor_inputs.keys()}""" + # ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError("`prompt` has to be of type `str` or `list[str]` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: " + f"{negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if max_sequence_length is not None and max_sequence_length > 256: + raise ValueError(f"`max_sequence_length` has to be less than or equal to 256 but is {max_sequence_length}") + + @staticmethod + def _prepare_latent_image_ids(batch_size, height, width, device, dtype): + latent_image_ids = torch.zeros(height, width, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids.reshape( + latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + + return latent_image_ids.to(device=device, dtype=dtype) + + @staticmethod + def _pack_latents(latents, batch_size, num_channel_latents, height, width): + latents = latents.view(batch_size, num_channel_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channel_latents * 4) + return latents + + @staticmethod + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + # 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 = int(2 * (int(height) // (vae_scale_factor * 2))) + width = int(2 * (int(width) // (vae_scale_factor * 2))) + + latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + latents = latents.reshape(batch_size, channels // (2 * 2), height, width) + return latents + + def prepare_latents( + self, + batch_size, + num_channel_latents, + height, + width, + dtype, + device, + 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 = int(2 * (int(height) // (self.vae_scale_factor * 2))) + width = int(2 * (int(width) // (self.vae_scale_factor * 2))) + + shape = (batch_size, num_channel_latents, height, width) + + if latents is not None: + latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + return latents.to(device=device, dtype=dtype), latent_image_ids + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents(latents, batch_size, num_channel_latents, height, width) + + latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + + return latents, latent_image_ids + + 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 + if hasattr(self.scheduler.config, "use_flow_sigmas") and self.scheduler.config.use_flow_sigmas: + sigmas = None + + 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._execution_device, + sigmas=sigmas, + mu=mu, + ) + return timesteps, num_inference_steps + + def denoising( + self, + latents, + timesteps, + prompt_embeds, + negative_prompt_embeds, + guidance_scale, + do_classifier_free_guidance, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + text_ids, + negative_text_ids, + latent_image_ids, + ): + self.scheduler.set_begin_index(0) + for i, t in enumerate(timesteps): + if self._interrupt: + break + + self._current_timestep = t + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + return_dict=False, + )[0] + + if do_classifier_free_guidance: + neg_noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + encoder_hidden_states=negative_prompt_embeds, + txt_ids=negative_text_ids, + img_ids=latent_image_ids, + return_dict=False, + )[0] + + noise_pred = neg_noise_pred + guidance_scale * (noise_pred - neg_noise_pred) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + # Not used in this pipeline + # if callback_on_step_end is not None: + # callback_kwargs = {} + # for k in callback_on_step_end_tensor_inputs: + # callback_kwargs[k] = locals()[k] + # callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + # latents = callback_outputs.pop("latents", latents) + # prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + guidance_scale: float = 5.0, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 256, + ) -> DiffusionOutput: + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + not greater than `1`). + guidance_scale (`float`, *optional*, defaults to 1.0): + True classifier-free guidance (guidance scale) is enabled when `guidance_scale` > 1 and + `negative_prompt` is provided. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`list[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`list`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.ovis_image.OvisImagePipelineOutput`] or `tuple`: + [`~pipelines.ovis_image.OvisImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is a list with the generated images. + """ + prompt = req.prompt if req.prompt is not None else prompt + negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt + height = req.height or self.default_sample_size * self.vae_scale_factor + width = req.width or self.default_sample_size * self.vae_scale_factor + num_inference_steps = req.num_inference_steps or num_inference_steps + guidance_scale = req.guidance_scale if req.guidance_scale is not None else guidance_scale + generator = req.generator or generator + req_num_outputs = getattr(req, "num_outputs_per_prompt", None) + if req_num_outputs and req_num_outputs > 0: + num_images_per_prompt = req_num_outputs + + # Steps: + # 1. Check Inputs + # 2. encode prompts + # 4. Prepare latents + # 5. Prepare timesteps + # 6. diffusion latents + # 7. decode latents + # 8. post process outputs + + self.check_inputs( + prompt, + height, + width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + device = self._execution_device + device = self._execution_device + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._current_timestep = None + 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) + else: + batch_size = prompt_embeds.shape[0] + + do_classifier_free_guidance = guidance_scale > 1.0 + prompt_embeds, text_ids = self.encode_prompt( + prompt=prompt, + prompt_embeds=prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + ) + + negative_text_ids = None + if do_classifier_free_guidance: + negative_prompt_embeds, negative_text_ids = self.encode_prompt( + prompt=negative_prompt, + prompt_embeds=negative_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + ) + + # 4. Prepare latent variables + num_channel_latents = self.transformer.in_channels // 4 + latents, latent_image_ids = self.prepare_latents( + batch_size=batch_size * num_images_per_prompt, + num_channel_latents=num_channel_latents, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + + # 5. Prepare timesteps + + image_seq_len = latents.shape[1] + timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, image_seq_len) + + # num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + if self.joint_attention_kwargs is None: + self._joint_attention_kwargs = {} + + # 6. Denoising loop + + # We set the index here to remove DtoH sync, helpful especially during compilation + # Check out more details here: https://github.com/huggingface/diffusers/pull/11696 + + latents = self.denoising( + latents, + timesteps, + prompt_embeds, + negative_prompt_embeds, + guidance_scale, + do_classifier_free_guidance, + callback_on_step_end, + callback_on_step_end_tensor_inputs, + text_ids, + negative_text_ids, + latent_image_ids, + ) + + self._current_timestep = None + if output_type == "latent": + image = latents + else: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + 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) + + 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 3c4cf3a1a..96357d6a5 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -29,6 +29,11 @@ _DIFFUSION_MODELS = { "pipeline_z_image", "ZImagePipeline", ), + "OvisImagePipeline": ( + "ovis_image", + "pipeline_ovis_image", + "OvisImagePipeline", + ), "WanPipeline": ( "wan2_2", "pipeline_wan2_2", @@ -77,6 +82,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func", "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func", "ZImagePipeline": "get_post_process_func", + "OvisImagePipeline": "get_ovis_image_post_process_func", "WanPipeline": "get_wan22_post_process_func", "LongCatImagePipeline": "get_longcat_image_post_process_func", }