diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b4a15f9de..30830f65f 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -40,6 +40,7 @@ th { |`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | |`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | +|`FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | |`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 4b5f6bf84..dbae38e86 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -41,6 +41,7 @@ The following table shows which models are currently supported by parallelism me | **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ | | **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ | +| **FLUX.1-Kontext-dev** | `black-forest-labs/FLUX.1-Kontext-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ | | **OmniGen2** | `OmniGen2/OmniGen2` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ❌ | !!! note "TP Limitations for Diffusion Models" diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index ad97909c5..88ea4e5e8 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -71,6 +71,7 @@ The following table shows which models are currently supported by each accelerat | **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | | **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | +| **FLUX.1-Kontext-dev** | `black-forest-labs/FLUX.1-Kontext-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ### VideoGen diff --git a/tests/e2e/offline_inference/test_flux_kontext.py b/tests/e2e/offline_inference/test_flux_kontext.py new file mode 100644 index 000000000..93dca21c9 --- /dev/null +++ b/tests/e2e/offline_inference/test_flux_kontext.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Tests for FluxKontext model pipeline. + +FluxKontext is a text-to-image and image-to-image diffusion model that supports: +- Text-to-image generation +- Image editing with text guidance +""" + +import os +import sys +from pathlib import Path + +import pytest +from PIL import Image + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +MODEL = "black-forest-labs/FLUX.1-Kontext-dev" + + +@pytest.mark.core_model +@pytest.mark.diffusion +def test_flux_kontext_text_to_image(): + """Test FluxKontext text-to-image generation with real model.""" + omni = Omni( + model=MODEL, + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + ), + enable_cpu_offload=False, + ) + + try: + omni_outputs = list( + omni.generate( + prompts=["A photo of a cat sitting on a laptop"], + sampling_params_list=OmniDiffusionSamplingParams( + height=512, + width=512, + num_inference_steps=2, + seed=42, + ), + ) + ) + + assert len(omni_outputs) > 0 + output = omni_outputs[0] + images = None + if output.images: + images = output.images + elif hasattr(output, "request_output") and output.request_output: + for stage_out in output.request_output: + if hasattr(stage_out, "images") and stage_out.images: + images = stage_out.images + break + + assert images is not None + assert len(images) > 0 + assert isinstance(images[0], Image.Image) + assert images[0].size == (512, 512) + finally: + omni.close() + + +@pytest.mark.core_model +@pytest.mark.diffusion +def test_flux_kontext_image_edit(): + """Test FluxKontext image-to-image editing with real model.""" + from vllm.assets.image import ImageAsset + + input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") + omni = Omni( + model=MODEL, + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + ), + enable_cpu_offload=False, + ) + + try: + omni_outputs = list( + omni.generate( + prompts=[ + { + "prompt": "Transform this image into a Vincent van Gogh style painting", + "multi_modal_data": {"img2img": input_image}, + "modalities": ["img2img"], + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=512, + width=512, + num_inference_steps=2, + seed=42, + ), + ) + ) + + assert len(omni_outputs) > 0 + output = omni_outputs[0] + images = None + if output.images: + images = output.images + elif hasattr(output, "request_output") and output.request_output: + for stage_out in output.request_output: + if hasattr(stage_out, "images") and stage_out.images: + images = stage_out.images + break + + assert images is not None + assert len(images) > 0 + assert isinstance(images[0], Image.Image) + assert images[0].size == (512, 512) + finally: + omni.close() diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py new file mode 100644 index 000000000..c13e1e818 --- /dev/null +++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py @@ -0,0 +1,166 @@ +""" +Comprehensive tests of diffusion features that are available in online serving mode +and are supported by the FluxKontext model. +""" + +import pytest + +from tests.conftest import ( + OmniServer, + OmniServerParams, + OpenAIClientHandler, + dummy_messages_from_mix_data, + generate_synthetic_image, +) + +EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." +NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist" +MODEL = "black-forest-labs/FLUX.1-Kontext-dev" + + +def _get_diffusion_feature_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--tensor-parallel-size", + "2", + "--enable-cpu-offload", + ], + ), + id="parallel_001", + ), + ] + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases(MODEL), + indirect=True, +) +def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Test text-to-image generation with FluxKontext in regular end-user scenarios.""" + messages = dummy_messages_from_mix_data(content_text="A photo of a cat sitting on a laptop keyboard") + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases(MODEL), + indirect=True, +) +def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Test image editing with FluxKontext in regular end-user scenarios.""" + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + + messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "negative_prompt": NEGATIVE_PROMPT, + "true_cfg_scale": 4.0, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases(MODEL), + indirect=True, +) +def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Test image editing with FluxKontext without negative prompt.""" + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + + messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases(MODEL), + indirect=True, +) +def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Test high-resolution generation with FluxKontext.""" + messages = dummy_messages_from_mix_data(content_text="A beautiful landscape with mountains and a lake") + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 768, + "width": 1024, + "num_inference_steps": 2, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases(MODEL), + indirect=True, +) +def test_flux_kontext_multiple_outputs(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Test generating multiple outputs with FluxKontext.""" + messages = dummy_messages_from_mix_data(content_text="A photo of a cat sitting on a laptop keyboard") + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "num_outputs_per_prompt": 2, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) diff --git a/vllm_omni/diffusion/models/flux/__init__.py b/vllm_omni/diffusion/models/flux/__init__.py index cad3508a1..6b13c4d96 100644 --- a/vllm_omni/diffusion/models/flux/__init__.py +++ b/vllm_omni/diffusion/models/flux/__init__.py @@ -1,17 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""FLUX.1-dev diffusion model components.""" +"""FLUX diffusion model components.""" from vllm_omni.diffusion.models.flux.flux_transformer import ( + FluxKontextTransformer2DModel, FluxTransformer2DModel, ) from vllm_omni.diffusion.models.flux.pipeline_flux import ( FluxPipeline, get_flux_post_process_func, ) +from vllm_omni.diffusion.models.flux.pipeline_flux_kontext import ( + FluxKontextPipeline, + get_flux_kontext_post_process_func, +) __all__ = [ "FluxPipeline", + "FluxKontextPipeline", "FluxTransformer2DModel", + "FluxKontextTransformer2DModel", "get_flux_post_process_func", + "get_flux_kontext_post_process_func", ] diff --git a/vllm_omni/diffusion/models/flux/flux_pipeline_mixin.py b/vllm_omni/diffusion/models/flux/flux_pipeline_mixin.py new file mode 100644 index 000000000..8e1e6b5b3 --- /dev/null +++ b/vllm_omni/diffusion/models/flux/flux_pipeline_mixin.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flux Pipeline Mixin - Shared methods for Flux pipelines.""" + +import inspect + +import torch + + +class FluxPipelineMixin: + """Shared methods for Flux pipelines (FluxPipeline, FluxKontextPipeline).""" + + @staticmethod + 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 + + @staticmethod + def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, + ): + 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 + + @staticmethod + def _prepare_latent_image_ids(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_channels_latents, height, width): + latents = latents.view(batch_size, num_channels_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_channels_latents * 4) + return latents + + @staticmethod + def _unpack_latents(latents, height, width, vae_scale_factor): + if len(latents.shape) == 4: + return latents + batch_size, num_patches, channels = latents.shape + + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 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 // 4, height, width) + + return latents + + @staticmethod + def retrieve_latents( + encoder_output: torch.Tensor, + generator: torch.Generator | None = None, + sample_mode: str = "sample", + ): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise ValueError("Encoder output does not have latent_dist or latents attribute.") diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 1a003b86c..8f81f3cd2 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -4,6 +4,7 @@ from collections.abc import Iterable from typing import TYPE_CHECKING, Any import torch +import torch.nn as nn import torch.nn.functional as F from diffusers.models.embeddings import ( CombinedTimestepGuidanceTextProjEmbeddings, @@ -13,7 +14,6 @@ from diffusers.models.embeddings import ( from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle from diffusers.utils import is_torch_npu_available -from torch import nn from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm @@ -25,6 +25,8 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -190,7 +192,8 @@ class FluxAttention(torch.nn.Module): self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, - image_rotary_emb: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: qkv, _ = self.to_qkv(hidden_states) @@ -226,10 +229,17 @@ class FluxAttention(torch.nn.Module): query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb) # [S, D/2] + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + hidden_states = self.attn( query, key, value, + attn_metadata, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -244,7 +254,6 @@ class FluxAttention(torch.nn.Module): return hidden_states, encoder_hidden_states else: - # For single-stream blocks, there's no to_out (RowParallelLinear) to handle the reduction if get_tensor_model_parallel_world_size() > 1: hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1) return hidden_states @@ -483,7 +492,7 @@ class FluxTransformer2DModel(nn.Module): def __init__( self, - od_config: OmniDiffusionConfig, + od_config: OmniDiffusionConfig = None, patch_size: int = 1, in_channels: int = 64, out_channels: int = None, @@ -495,18 +504,23 @@ class FluxTransformer2DModel(nn.Module): pooled_projection_dim: int = 768, guidance_embeds: bool = True, axes_dims_rope: tuple[int, int, int] = (16, 56, 56), + theta: float = 10000.0, quant_config: "QuantizationConfig | None" = None, ): super().__init__() - model_config = od_config.tf_model_config - num_layers = model_config.num_layers - self.parallel_config = od_config.parallel_config + if od_config is not None: + model_config = od_config.tf_model_config + num_layers = model_config.num_layers + self.parallel_config = od_config.parallel_config + else: + self.parallel_config = None + self.in_channels = in_channels self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.guidance_embeds = guidance_embeds - self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope) + self.pos_embed = FluxPosEmbed(theta=theta, axes_dim=axes_dims_rope) text_time_guidance_cls = ( CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings ) @@ -687,3 +701,102 @@ class FluxTransformer2DModel(nn.Module): loaded_params.add(original_name) loaded_params.add(lookup_name) return loaded_params + + +class FluxKontextTransformer2DModel(FluxTransformer2DModel): + def __init__( + self, + od_config: OmniDiffusionConfig = None, + patch_size: int = 1, + in_channels: int = 64, + out_channels: int = None, + num_layers: int = 19, + num_single_layers: int = 38, + attention_head_dim: int = 128, + num_attention_heads: int = 24, + joint_attention_dim: int = 4096, + pooled_projection_dim: int = 768, + guidance_embeds: bool = True, + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), + theta: float = 10000.0, + ): + super().__init__( + od_config=od_config, + patch_size=patch_size, + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + num_single_layers=num_single_layers, + attention_head_dim=attention_head_dim, + num_attention_heads=num_attention_heads, + joint_attention_dim=joint_attention_dim, + pooled_projection_dim=pooled_projection_dim, + guidance_embeds=guidance_embeds, + axes_dims_rope=axes_dims_rope, + theta=theta, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + pooled_projections: torch.Tensor = None, + timestep: torch.LongTensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + guidance: torch.Tensor | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + hidden_states = self.x_embedder(hidden_states) + timestep = timestep.to(hidden_states.dtype) * 1000 + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + + if guidance is None: + temb = self.time_text_embed(timestep, pooled_projections) + else: + temb = self.time_text_embed(timestep, guidance, pooled_projections) + + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + if txt_ids is not None and 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 is not None and 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) + image_rotary_emb = self.pos_embed(ids) + + for block in 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, + joint_attention_kwargs=joint_attention_kwargs, + ) + + for block in 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, + joint_attention_kwargs=joint_attention_kwargs, + ) + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index b7aa0c7b4..3dde14866 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import inspect import json import logging import os @@ -27,6 +26,7 @@ from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_g 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.flux import FluxTransformer2DModel +from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers @@ -48,6 +48,11 @@ def get_flux_post_process_func( 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 @@ -60,80 +65,7 @@ def get_flux_post_process_func( 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 - - -# 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, -) -> tuple[torch.Tensor, int]: - 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: - `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: - 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 FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): +class FluxPipeline(nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -401,44 +333,6 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): return prompt_embeds, pooled_prompt_embeds, text_ids - @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_channels_latents, height, width): - latents = latents.view(batch_size, num_channels_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_channels_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 = 2 * (int(height) // (vae_scale_factor * 2)) - width = 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, @@ -458,7 +352,7 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): shape = (batch_size, num_channels_latents, height, width) if latents is not None: - latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + latent_image_ids = self._prepare_latent_image_ids(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: @@ -470,20 +364,20 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) - latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + latent_image_ids = self._prepare_latent_image_ids(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 - mu = calculate_shift( + mu = self.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( + timesteps, num_inference_steps = self.retrieve_timesteps( self.scheduler, num_inference_steps, sigmas=sigmas, diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py new file mode 100644 index 000000000..bf741a64e --- /dev/null +++ b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py @@ -0,0 +1,699 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# This implementation is heavily inspired by the diffusers project. +# Original implementation: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux_kontext.py + +from __future__ import annotations + +import json +import os +from collections.abc import Callable, Iterable +from typing import Any, cast + +import numpy as np +import PIL.Image +import torch +import torch.nn as nn +from diffusers.image_processor import VaeImageProcessor +from diffusers.loaders import TextualInversionLoaderMixin +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils.torch_utils import randn_tensor +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast +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.flux import ( + FluxKontextTransformer2DModel, +) +from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin +from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.logger import init_logger + +logger = init_logger(__name__) + + +def get_flux_kontext_post_process_func(od_config: OmniDiffusionConfig) -> Callable: + """Get postprocessing function for FLUX.1-Kontext pipeline.""" + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + + 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) -> list[PIL.Image.Image]: + result = image_processor.postprocess(images) + return result + + return post_process_func + + +class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput): + """FLUX.1-Kontext pipeline for image editing with text guidance.""" + + support_image_input = True + + 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 = CLIPTextModel.from_pretrained( + model, + subfolder="text_encoder", + local_files_only=local_files_only, + ).to(self._execution_device) + + self.text_encoder_2 = T5EncoderModel.from_pretrained( + model, + subfolder="text_encoder_2", + local_files_only=local_files_only, + ).to(self._execution_device) + + self.tokenizer = CLIPTokenizer.from_pretrained( + model, + subfolder="tokenizer", + local_files_only=local_files_only, + ) + self.tokenizer_2 = T5TokenizerFast.from_pretrained( + model, + subfolder="tokenizer_2", + local_files_only=local_files_only, + ) + + self.vae = AutoencoderKL.from_pretrained( + model, + subfolder="vae", + local_files_only=local_files_only, + ).to(self._execution_device) + + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, FluxKontextTransformer2DModel) + self.transformer = FluxKontextTransformer2DModel(**transformer_kwargs) + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.tokenizer_max_length = ( + self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 + ) + self.default_sample_size = 128 + self._guidance_scale = None + self._attention_kwargs = None + self._joint_attention_kwargs = None + self._num_timesteps = None + self._interrupt = False + self._callback_tensor_inputs = ["latents", "prompt_embeds"] + self.latent_channels = self.vae.config.latent_channels if hasattr(self.vae, "config") else 16 + + def _get_t5_prompt_embeds( + self, + prompt: str | list[str] = None, + num_images_per_prompt: int = 1, + max_sequence_length: int = 512, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, self.tokenizer_2) + + text_inputs = self.tokenizer_2( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_length=False, + return_overflowing_tokens=False, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {max_sequence_length} tokens: {removed_text}" + ) + + prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0] + + dtype = self.text_encoder_2.dtype + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + _, 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 _get_clip_prompt_embeds( + self, + prompt: str | list[str], + num_images_per_prompt: int = 1, + device: torch.device | None = None, + ): + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, self.tokenizer) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer_max_length, + truncation=True, + return_overflowing_tokens=False, + return_length=False, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + untruncated_ids = self.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 = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer_max_length} tokens: {removed_text}" + ) + prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) + + prompt_embeds = prompt_embeds.pooler_output + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) + + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + return prompt_embeds + + def encode_prompt( + self, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + max_sequence_length: int = 512, + ): + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt_embeds is None: + prompt_2 = prompt_2 or prompt + prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 + + # We only use the pooled prompt output from the CLIPTextModel + pooled_prompt_embeds = self._get_clip_prompt_embeds( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + ) + prompt_embeds = self._get_t5_prompt_embeds( + prompt=prompt_2, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + 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).to(device=device, dtype=dtype) + + return prompt_embeds, pooled_prompt_embeds, text_ids + + def check_inputs( + self, + prompt, + prompt_2=None, + height=None, + width=None, + negative_prompt=None, + negative_prompt_2=None, + prompt_embeds=None, + pooled_prompt_embeds=None, + negative_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + max_sequence_length=None, + ): + if height is not None and height % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by " + f"{self.vae_scale_factor * 2} but are {height} and {width}. " + f"Dimensions 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 + ): + invalid_keys = [k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs] + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in " + f"{self._callback_tensor_inputs}, but found {invalid_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_2 is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} 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(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)}") + + 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." + ) + elif negative_prompt_2 is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and pooled_prompt_embeds is None: + raise ValueError( + "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. " + "Make sure to generate `pooled_prompt_embeds` from the same text encoder " + "that was used to generate `prompt_embeds`." + ) + if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` " + "also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` " + "from the same text encoder that was used to generate `negative_prompt_embeds`." + ) + + if max_sequence_length is not None and max_sequence_length > 512: + raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}") + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if image.ndim != 4: + raise ValueError(f"Expected image dims 4, got {image.ndim}.") + + if isinstance(generator, list): + image_latents = [ + self.retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax") + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = self.retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") + + image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + + if image_latents.ndim == 5: + logger.warning("VAE output has 5 dimensions, squeezing first dimension") + image_latents = image_latents.squeeze(0) + + return image_latents + + def prepare_latents( + self, + image: torch.Tensor | None, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ): + 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." + ) + + 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) + + image_latents = image_ids = None + if image is not None: + image = image.to(device=device, dtype=dtype) + if image.shape[1] != self.latent_channels: + image_latents = self._encode_vae_image(image=image, generator=generator) + else: + image_latents = image + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + # expand init_latents for batch_size + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + else: + image_latents = torch.cat([image_latents], dim=0) + + image_latent_height, image_latent_width = image_latents.shape[2:] + image_latents = self._pack_latents( + image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width + ) + image_ids = self._prepare_latent_image_ids(image_latent_height // 2, image_latent_width // 2, device, dtype) + # image ids are the same as latent ids with the first dimension set to 1 instead of 0 + image_ids[..., 0] = 1 + + latent_ids = self._prepare_latent_image_ids(height // 2, width // 2, device, dtype) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + else: + latents = latents.to(device=device, dtype=dtype) + if latents.dim() == 4: + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + + return latents, image_latents, latent_ids, image_ids + + @property + def interrupt(self): + return self._interrupt + + @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 + + @torch.no_grad() + def forward( + self, + req: OmniDiffusionRequest, + image: PIL.Image.Image | list[PIL.Image.Image] | None = None, + prompt: str | list[str] | None = None, + prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 28, + guidance_scale: float = 3.5, + true_cfg_scale: float = 1.0, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | 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, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + sigmas: list[float] | None = None, + ) -> DiffusionOutput: + # Handle multiple prompts - only take the first one, similar to Flux2KleinPipeline + if len(req.prompts) > 1: + logger.warning( + "This model only supports a single prompt, not a batched request. Taking only the first prompt for now." + ) + first_prompt = req.prompts[0] if req.prompts else None + + if isinstance(first_prompt, str): + prompt = first_prompt + elif first_prompt: + prompt = first_prompt.get("prompt") or "" + # else: keep original prompt unchanged + + # Handle image from prompt data + if isinstance(first_prompt, str): + raw_image = None + elif first_prompt: + raw_image = first_prompt.get("multi_modal_data", {}).get("image") + else: + raw_image = None + + if raw_image is None: + pass # use image from param list + elif isinstance(raw_image, list): + image = [PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image, im) for im in raw_image] + else: + image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image) + height = req.sampling_params.height or height + width = req.sampling_params.width or width + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + guidance_scale = req.sampling_params.guidance_scale or guidance_scale + generator = req.sampling_params.generator or generator + latents = ( + req.sampling_params.extra_args.get("latents") + if req.sampling_params.extra_args.get("latents") is not None + else latents + ) + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + self.check_inputs( + prompt, + prompt_2=prompt_2, + height=height, + width=width, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._joint_attention_kwargs = attention_kwargs if attention_kwargs is not None else {} + self._interrupt = False + + 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] + + device = self._execution_device + + 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.0 and has_neg_prompt + + # 1. Prepare text embeddings + prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + device=device, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + max_sequence_length=max_sequence_length, + ) + + if do_true_cfg: + negative_prompt_embeds, negative_pooled_prompt_embeds, negative_text_ids = self.encode_prompt( + prompt=negative_prompt, + prompt_2=negative_prompt_2, + prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=negative_pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + + # 2. Process images + if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): + img = image[0] if isinstance(image, list) else image + image_height, image_width = self.image_processor.get_default_height_width(img) + multiple_of = self.vae_scale_factor * 2 + image_width = image_width // multiple_of * multiple_of + image_height = image_height // multiple_of * multiple_of + image = self.image_processor.resize(image, image_height, image_width) + image = self.image_processor.preprocess(image, image_height, image_width) + if height is None and width is None: + height = image_height + width = image_width + + # 3. Prepare latent variables + num_channels_latents = self.transformer.in_channels // 4 + latents, image_latents, latent_ids, image_ids = self.prepare_latents( + image=image, + batch_size=batch_size * num_images_per_prompt, + num_channels_latents=num_channels_latents, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + + if image_ids is not None: + latent_ids = torch.cat([latent_ids, image_ids], dim=0) + + # 4. Prepare timesteps + 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 + image_seq_len = latents.shape[1] + mu = self.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 = self.retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + self._num_timesteps = len(timesteps) + + if self.transformer.guidance_embeds: + guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32) + guidance = guidance.expand(latents.shape[0]) + else: + guidance = None + + if self.joint_attention_kwargs is None: + self._joint_attention_kwargs = {} + + # 5. Denoising loop + self.scheduler.set_begin_index(0) + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1) + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + noise_pred = noise_pred[:, : latents.size(1)] + + if do_true_cfg: + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=negative_pooled_prompt_embeds, + encoder_hidden_states=negative_prompt_embeds, + txt_ids=negative_text_ids, + img_ids=latent_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + neg_noise_pred = neg_noise_pred[:, : latents.size(1)] + noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) + + 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(): + latents = latents.to(latents_dtype) + + 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) + 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 d4a738b81..994dac04a 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -97,6 +97,11 @@ _DIFFUSION_MODELS = { "pipeline_sd3", "StableDiffusion3Pipeline", ), + "FluxKontextPipeline": ( + "flux", + "pipeline_flux_kontext", + "FluxKontextPipeline", + ), "HunyuanImage3ForCausalMM": ( "hunyuan_image_3", "pipeline_hunyuan_image_3", @@ -319,6 +324,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "BagelPipeline": "get_bagel_post_process_func", "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", + "FluxKontextPipeline": "get_flux_kontext_post_process_func", "Flux2KleinPipeline": "get_flux2_klein_post_process_func", "NextStep11Pipeline": "get_nextstep11_post_process_func", "FluxPipeline": "get_flux_post_process_func",