diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 7ff13fd1b..6703dd9c7 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -382,7 +382,7 @@ def forward( for p in req.prompts ] - # Extract sampling parameters + # Extract common sampling parameters sampling_params = req.sampling_params num_inference_steps = sampling_params.num_inference_steps or 50 guidance_scale = sampling_params.guidance_scale or 7.5 diff --git a/docs/design/feature/cfg_parallel.md b/docs/design/feature/cfg_parallel.md index c73a87749..053360494 100644 --- a/docs/design/feature/cfg_parallel.md +++ b/docs/design/feature/cfg_parallel.md @@ -126,16 +126,16 @@ Call `self.diffuse` in your pipeline's forward function: ```python import torch.nn as nn +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + class YourModelPipeline(nn.Module, CFGParallelMixin): def forward( self, - prompt: str, - negative_prompt: str | None = None, - guidance_scale: float = 3.5, - num_inference_steps: int = 50, - **kwargs, - ): - # Encode prompts, Initialize latents, Get timesteps + req: DiffusionRequestBatch, + ) -> list[DiffusionOutput]: + # Read prompts and sampling params from the request + ... + # Encode prompts, initialize latents, get timesteps ... # Run diffusion loop (calls the mixin's diffuse method) latents = self.diffuse( diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md index dc8b645e4..31334526a 100644 --- a/docs/features/custom_pipeline.md +++ b/docs/features/custom_pipeline.md @@ -54,20 +54,24 @@ Create a custom pipeline class that extends an existing pipeline. In this exampl ```python # custom_pipeline.py -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch import torch class CustomPipeline(QwenImageEditPipeline): def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) - def forward(self, req, prompt=None, negative_prompt=None, **kwargs): + def forward(self, req: DiffusionRequestBatch): + # Optionally customize sampling parameters + actual_num_steps = req.sampling_params.num_inference_steps or 50 + req.sampling_params.num_inference_steps = actual_num_steps + # Call parent's forward to get normal output - output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs) + output = super().forward(req=req) # Add custom trajectory data - actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50) output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32) output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32) diff --git a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py index b1c74e2b3..fb14d9d33 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py @@ -2,14 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging -from typing import Any -import PIL.Image import torch -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch logger = logging.getLogger(__name__) @@ -18,58 +16,14 @@ class CustomPipeline(QwenImageEditPipeline): def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - image: PIL.Image.Image | torch.Tensor | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_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, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch): """Forward pass for image editing with dummy trajectory data.""" - # Call parent's forward to get the normal output - output = super().forward( - req=req, - prompt=prompt, - negative_prompt=negative_prompt, - image=image, - true_cfg_scale=true_cfg_scale, - height=height, - width=width, - num_inference_steps=num_inference_steps, - sigmas=sigmas, - guidance_scale=guidance_scale, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - prompt_embeds=prompt_embeds, - prompt_embeds_mask=prompt_embeds_mask, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_mask=negative_prompt_embeds_mask, - output_type=output_type, - attention_kwargs=attention_kwargs, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) + # Customize num_inference_steps that best suits your model + actual_num_steps = req.sampling_params.num_inference_steps or 50 + req.sampling_params.num_inference_steps = actual_num_steps - # Get actual num_inference_steps used - actual_num_steps = req.sampling_params.num_inference_steps or num_inference_steps + # Call parent's forward to get the normal output + output = super().forward(req=req) # Create dummy trajectory data dummy_trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32) diff --git a/tests/diffusion/models/flux/test_flux_pipeline.py b/tests/diffusion/models/flux/test_flux_pipeline.py index fd5d7a2ac..5a06133f4 100644 --- a/tests/diffusion/models/flux/test_flux_pipeline.py +++ b/tests/diffusion/models/flux/test_flux_pipeline.py @@ -21,6 +21,8 @@ def _make_flux_sampling(**overrides): "true_cfg_scale": 4.0, "num_outputs_per_prompt": 0, "latents": None, + "output_type": None, + "max_sequence_length": None, } values.update(overrides) return SimpleNamespace(**values) @@ -131,7 +133,7 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch): "negative_prompt_embeds": negative_prompt_embeds_a, "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds_a, }, - sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a), + sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a, output_type="latent"), ), SimpleNamespace( request_id="flux-prompt-b", @@ -145,12 +147,12 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch): "negative_pooled_prompt_embeds": [negative_pooled_prompt_embeds_b], }, }, - sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b), + sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b, output_type="latent"), ), ] ) - outputs = pipeline.forward(batch, output_type="latent") + outputs = pipeline.forward(batch) assert encode_calls[0]["prompt"] is None assert encode_calls[0]["prompt_2"] is None diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py index 3f85a6b17..7f23b4733 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py @@ -1,4 +1,3 @@ -import inspect from types import SimpleNamespace import pytest @@ -193,6 +192,7 @@ def _make_request_batch_prompt_sampling(**overrides): "true_cfg_scale": None, "guidance_scale_provided": False, "guidance_scale": 1.0, + "output_type": None, } values.update(overrides) return SimpleNamespace(**values) @@ -344,16 +344,3 @@ def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_cl with pytest.raises(AssertionError, match="text encoder should not run"): pipeline.encode_prompt(prompt="short prompt") - - -@pytest.mark.parametrize( - "pipeline_class", - [ - QwenImagePipeline, - QwenImageLayeredPipeline, - QwenImageEditPipeline, - QwenImageEditPlusPipeline, - ], -) -def test_forward_max_sequence_length_default_is_1024(pipeline_class: type): - assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024 diff --git a/tests/diffusion/models/sd3/test_sd3_pipeline.py b/tests/diffusion/models/sd3/test_sd3_pipeline.py index 3164e5dda..4b0bf843f 100644 --- a/tests/diffusion/models/sd3/test_sd3_pipeline.py +++ b/tests/diffusion/models/sd3/test_sd3_pipeline.py @@ -21,6 +21,7 @@ def _make_sd3_sampling(**overrides): "generator": None, "latents": None, "guidance_scale": 4.0, + "output_type": None, } values.update(overrides) return SimpleNamespace(**values) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index d695766f6..198c31908 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -20,6 +20,12 @@ class _StubTransformer(nn.Module): return torch.float32 +class _StubTextEncoder(nn.Module): + @property + def dtype(self) -> torch.dtype: + return torch.float32 + + class _StubScheduler: def __init__(self, timesteps: list[int]) -> None: self.timesteps = torch.tensor(timesteps, dtype=torch.int64) @@ -41,12 +47,30 @@ def _noop_progress_bar(*args, **kwargs): yield _Bar() +def _stub_encode_prompt( + prompt, + negative_prompt=None, + do_classifier_free_guidance=True, + num_videos_per_prompt=1, + max_sequence_length=512, + device=None, + dtype=None, +): + del negative_prompt, do_classifier_free_guidance, device, dtype + batch_size = 1 if isinstance(prompt, str) else len(prompt) + n = batch_size * num_videos_per_prompt + hidden_size = 8 + prompt_embeds = torch.zeros(n, max_sequence_length, hidden_size) + return prompt_embeds, None + + def _make_pipeline() -> Wan22Pipeline: pipeline = object.__new__(Wan22Pipeline) nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") pipeline.transformer = _StubTransformer() pipeline.transformer_2 = None + pipeline.text_encoder = _StubTextEncoder() pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4) pipeline.scheduler = _StubScheduler([9, 5]) pipeline.od_config = SimpleNamespace(flow_shift=5.0) @@ -61,15 +85,14 @@ def _make_pipeline() -> Wan22Pipeline: pipeline._num_timesteps = None pipeline._current_timestep = None pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = _stub_encode_prompt # type: ignore[method-assign] pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32) pipeline.progress_bar = _noop_progress_bar return pipeline -def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None: +def test_forward_delegates_denoising_to_diffuse() -> None: pipeline = _make_pipeline() - - prompt_embeds = torch.randn(1, 8) captured: dict[str, object] = {} def _fake_diffuse(**kwargs): @@ -86,8 +109,8 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None: width=None, num_frames=1, num_inference_steps=2, - guidance_scale_provided=False, - guidance_scale=None, + guidance_scale_provided=True, + guidance_scale=1.0, guidance_scale_2=None, boundary_ratio=None, generator=None, @@ -95,14 +118,16 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None: num_outputs_per_prompt=1, max_sequence_length=32, latents=None, + output_type="latent", extra_args={}, ), ) batch = DiffusionRequestBatch(requests=[mock_req]) - output = pipeline.forward(batch, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0) + output = pipeline.forward(batch) assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8))) + assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8)) assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps) assert captured["guidance_low"] == 1.0 assert captured["guidance_high"] == 1.0 diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index bec8d3a14..e7f97597b 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -502,38 +502,14 @@ class FluxPipeline( return False return True - def forward( - self, - req: DiffusionRequestBatch, - 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, - true_cfg_scale: float = 1.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 28, - sigmas: list[float] | None = None, - guidance_scale: float = 3.5, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: torch.FloatTensor | None = None, - pooled_prompt_embeds: torch.FloatTensor | None = None, - negative_prompt_embeds: torch.FloatTensor | None = None, - negative_pooled_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_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, - ) -> list[DiffusionOutput]: + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: """Forward pass for flux.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. sampling_params_list = req.sampling_params_list common_sampling_params = sampling_params_list[0] - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: @@ -541,40 +517,39 @@ class FluxPipeline( height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor - num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps - sigmas = common_sampling_params.sigmas or sigmas - guidance_scale = ( - common_sampling_params.guidance_scale - if common_sampling_params.guidance_scale is not None - else guidance_scale - ) - true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale + num_inference_steps = common_sampling_params.num_inference_steps or 28 + sigmas = common_sampling_params.sigmas + guidance_scale = common_sampling_params.guidance_scale + true_cfg_scale = common_sampling_params.true_cfg_scale or 1.0 num_images_per_prompt = ( - common_sampling_params.num_outputs_per_prompt - if common_sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1 ) - generator = req.collate_request_generators(num_images_per_prompt, generator) - latents = req.collate_request_tensors("latents", latents) + max_sequence_length = common_sampling_params.max_sequence_length or 512 + generator = req.collate_request_generators(num_images_per_prompt, None) + latents = req.collate_request_tensors("latents", None) prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( req.prompts, { - "prompt_embeds": prompt_embeds, - "negative_prompt_embeds": negative_prompt_embeds, - "pooled_prompt_embeds": pooled_prompt_embeds, - "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds, + "prompt_embeds": None, + "negative_prompt_embeds": None, + "pooled_prompt_embeds": None, + "negative_pooled_prompt_embeds": None, }, ) prompt_embeds = prompt_fields["prompt_embeds"] negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] pooled_prompt_embeds = prompt_fields["pooled_prompt_embeds"] negative_pooled_prompt_embeds = prompt_fields["negative_pooled_prompt_embeds"] + prompt_2 = None + negative_prompt_2 = None if prompt_embeds is not None: prompt = None - prompt_2 = None if negative_prompt_embeds is not None: negative_prompt = None - negative_prompt_2 = None + + output_type = common_sampling_params.output_type or "pil" + joint_attention_kwargs: dict[str, Any] | None = None + callback_on_step_end_tensor_inputs = ["latents"] # 1. Check inputs. Raise error if not correct self.check_inputs( @@ -598,9 +573,7 @@ class FluxPipeline( 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): + if prompt is not None: batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 81a6f50f1..184b5d402 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -750,108 +750,7 @@ class Flux2KleinPipeline( def interrupt(self): return self._interrupt - def forward( - self, - req: DiffusionRequestBatch, - image: PIL.Image.Image | list[PIL.Image.Image] | None = None, - reference_image: PIL.Image.Image | list[PIL.Image.Image] | None = None, - mask_image: PIL.Image.Image | list[PIL.Image.Image] | None = None, - prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - strength: float = 1.0, - guidance_scale: float | None = 4.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, - negative_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, - text_encoder_out_layers: tuple[int, ...] = (9, 18, 27), - padding_mask_crop: int | None = None, - ) -> DiffusionOutput: - r""" - Function invoked when calling the pipeline for generation. - - Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list of these): - `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both - numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list - or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a - list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image - latents as `image`, but if passing latents directly it is not encoded again. - 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. - guidance_scale (`float`, *optional*, defaults to 4.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. - of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting - `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to - the text `prompt`, usually at the expense of lower image quality. For step-wise distilled models, - `guidance_scale` is ignored. - 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.Tensor`, *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.Tensor`, *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.Tensor`, *optional*): - Pre-generated negative text embeddings. Note that "" is used as the negative prompt in this pipeline. - If not provided, will be generated from "". - 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.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple. - 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`. - text_encoder_out_layers (`Tuple[int]`): - Layer indices to use in the `text_encoder` to derive the final prompt embeddings. - - Examples: - - Returns: - [`~pipelines.flux2.Flux2PipelineOutput`] or `tuple`: [`~pipelines.flux2.Flux2PipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the - generated images. - """ + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: if len(req.prompts) > 1: logger.warning( """This model only supports a single prompt, not a batched request.""", @@ -878,39 +777,32 @@ class Flux2KleinPipeline( 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 + prompt_embeds = None + negative_prompt_embeds = None + + height = req.sampling_params.height + width = req.sampling_params.width num_inference_steps = ( - req.sampling_params.num_inference_steps - if req.sampling_params.num_inference_steps is not None - else num_inference_steps + 50 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps ) - sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) - generator = req.sampling_params.generator or generator + sigmas = req.sampling_params.sigmas + guidance_scale = req.sampling_params.guidance_scale + generator = req.sampling_params.generator num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers) + max_sequence_length = req.sampling_params.max_sequence_length or 512 + latents = req.sampling_params.latents + strength = req.sampling_params.strength if req.sampling_params.strength is not None else 1.0 - req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts] - if any(p is not None for p in req_prompt_embeds): - # If at list one prompt is provided as an embedding, - # Then assume that the user wants to provide embeddings for all prompts, and enter this if block - # If the user in fact provides mixed input format, req_prompt_embeds will have some None's - # And `torch.stack` automatically raises an exception for us - prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore # intentionally expect TypeError - - req_negative_prompt_embeds = [ - p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts - ] - if any(p is not None for p in req_negative_prompt_embeds): - negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore # intentionally expect TypeError + text_encoder_out_layers: tuple[int, ...] = req.sampling_params.extra_args.get( + "text_encoder_out_layers", (9, 18, 27) + ) + padding_mask_crop: int | None = req.sampling_params.extra_args.get("padding_mask_crop") + attention_kwargs: dict[str, Any] | None = None + callback_on_step_end: Callable[[int, int, dict], None] | None = None + callback_on_step_end_tensor_inputs = ["latents"] + output_type = req.sampling_params.output_type or "pil" # 1. Check inputs. Raise error if not correct self.check_inputs( @@ -930,12 +822,7 @@ class Flux2KleinPipeline( 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] + batch_size = 1 device = self._execution_device diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py index 2884adfca..b95f59228 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py @@ -495,66 +495,35 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile noise_pred = comb_pred * scale return noise_pred - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 4.5, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - joint_attention_kwargs: dict[str, Any] | None = None, - enable_cfg_renorm: bool | None = True, - cfg_renorm_min: float | None = 0.0, - enable_prompt_rewrite: bool | None = True, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - height = req.sampling_params.height or height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or width or self.default_sample_size * self.vae_scale_factor - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - generator = req.sampling_params.generator or generator - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) + prompt_embeds = None + negative_prompt_embeds = None + + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + generator = req.sampling_params.generator + guidance_scale = req.sampling_params.guidance_scale num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt is not None - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) - enable_prompt_rewrite = req.sampling_params.extra_args.get("enable_prompt_rewrite", enable_prompt_rewrite) - enable_cfg_renorm = req.sampling_params.extra_args.get("enable_cfg_renorm", enable_cfg_renorm) - cfg_renorm_min = req.sampling_params.extra_args.get("cfg_renorm_min", cfg_renorm_min) + latents = req.sampling_params.latents - req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts] - if any(p is not None for p in req_prompt_embeds): - # If at list one prompt is provided as an embedding, - # Then assume that the user wants to provide embeddings for all prompts, and enter this if block - # If the user in fact provides mixed input format, req_prompt_embeds will have some None's - # And `torch.stack` automatically raises an exception for us - prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore # intentionally expect TypeError - - req_negative_prompt_embeds = [ - p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts - ] - if any(p is not None for p in req_negative_prompt_embeds): - negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore # intentionally expect TypeError + extra_args = getattr(req.sampling_params, "extra_args", {}) or {} + enable_prompt_rewrite: bool = extra_args.get("enable_prompt_rewrite", True) + enable_cfg_renorm: bool = extra_args.get("enable_cfg_renorm", True) + cfg_renorm_min: float = extra_args.get("cfg_renorm_min", 0.0) + joint_attention_kwargs: dict[str, Any] | None = None + output_type = req.sampling_params.output_type or "pil" self.check_inputs( prompt, @@ -571,12 +540,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile 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] + batch_size = len(prompt) device = self.device if enable_prompt_rewrite and prompt is not None: diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py index 195aa89a4..4821425b0 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py @@ -6,7 +6,7 @@ import math import os import re from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -542,24 +542,7 @@ class LongCatImageEditPipeline( f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) - def forward( - self, - req: DiffusionRequestBatch, - image: PIL.Image.Image | torch.Tensor | None = None, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 3.5, - num_images_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - joint_attention_kwargs: dict[str, Any] | None = None, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. if len(req.prompts) > 1: @@ -570,27 +553,22 @@ class LongCatImageEditPipeline( first_prompt = req.prompts[0] prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") - prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("prompt_embeds") - negative_prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt_embeds") # type: ignore # Why it is list[torch.Tensor] in OmniTokenInputs or OmniEmbedsPrompt? Doesn't make sense - sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + sigmas = req.sampling_params.sigmas + guidance_scale = req.sampling_params.guidance_scale + num_inference_steps = req.sampling_params.num_inference_steps or 50 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt is not None - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) - generator = req.sampling_params.generator or generator + generator = req.sampling_params.generator height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + output_type = req.sampling_params.output_type or "pil" - if prompt is not None: - batch_size = 1 if isinstance(prompt, str) else len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = 1 if isinstance(prompt, str) else len(prompt) if not isinstance(first_prompt, str) and "preprocessed_image" in ( additional_information := first_prompt.get("additional_information", {}) @@ -600,13 +578,7 @@ class LongCatImageEditPipeline( calculated_height = additional_information.get("calculated_height", height) calculated_width = additional_information.get("calculated_width", width) else: - image_size = image[0].size if isinstance(image, list) else image.size - calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1]) - - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): - image = self.image_processor.resize(image, calculated_height, calculated_width) - prompt_image = self.image_processor.resize(image, calculated_height // 2, calculated_width // 2) - image = self.image_processor.preprocess(image, calculated_height, calculated_width) + raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.") self.check_inputs( prompt, @@ -621,7 +593,8 @@ class LongCatImageEditPipeline( prompt=prompt, image=prompt_image, prompt_embeds=prompt_embeds, num_images_per_prompt=num_images_per_prompt ) - if guidance_scale > 1: + do_true_cfg = guidance_scale > 1 and negative_prompt is not None + if do_true_cfg: (negative_prompt_embeds, negative_text_ids) = self.encode_prompt( prompt=negative_prompt, image=prompt_image, @@ -680,7 +653,6 @@ class LongCatImageEditPipeline( latent_model_input = torch.cat([latents, image_latents], dim=1) timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) - do_true_cfg = guidance_scale > 1 positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep / 1000, diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py index 371af3f34..0baad965b 100644 --- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py +++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py @@ -18,7 +18,7 @@ import inspect import json import os -from collections.abc import Callable, Iterable +from collections.abc import Iterable from typing import Any, ClassVar import numpy as np @@ -551,117 +551,33 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi def interrupt(self): return self._interrupt - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = 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. - """ + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] + prompt_embeds = None + negative_prompt_embeds = None + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) - generator = req.sampling_params.generator or generator + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + guidance_scale = req.sampling_params.guidance_scale + generator = req.sampling_params.generator num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + max_sequence_length = req.sampling_params.max_sequence_length or 256 + + joint_attention_kwargs: dict[str, Any] | None = None + output_type = req.sampling_params.output_type or "pil" + callback_on_step_end_tensor_inputs = ["latents"] # Steps: # 1. Check Inputs @@ -691,12 +607,7 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi 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] + batch_size = len(prompt) do_classifier_free_guidance = guidance_scale > 1.0 prompt_embeds, text_ids = self.encode_prompt( diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 43cd4bdc4..638316039 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -978,58 +978,34 @@ class QwenImagePipeline( height = state.sampling.height or self.default_sample_size * self.vae_scale_factor width = state.sampling.width or self.default_sample_size * self.vae_scale_factor - output_type = kwargs.get("output_type", "pil") + output_type = kwargs.get("output_type") or state.sampling.output_type or "pil" return self._decode_latents(state.latents, height, width, output_type) - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_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, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 1024, - ) -> list[DiffusionOutput]: + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: sampling_params_list = req.sampling_params_list common_sampling_params = sampling_params_list[0] extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) - prompt = extracted_prompt or prompt + prompt = extracted_prompt height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) - num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps - sigmas = common_sampling_params.sigmas or sigmas - max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length + num_inference_steps = common_sampling_params.num_inference_steps or 50 + sigmas = common_sampling_params.sigmas + max_sequence_length = common_sampling_params.max_sequence_length or 1024 num_images_per_prompt = ( - common_sampling_params.num_outputs_per_prompt - if common_sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1 ) - generator = req.collate_request_generators(num_images_per_prompt, generator) - latents = req.collate_request_tensors("latents", latents) + generator = req.collate_request_generators(num_images_per_prompt, None) + latents = req.collate_request_tensors("latents", None) prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( req.prompts, { - "prompt_embeds": prompt_embeds, - "prompt_embeds_mask": prompt_embeds_mask, - "negative_prompt_embeds": negative_prompt_embeds, - "negative_prompt_embeds_mask": negative_prompt_embeds_mask, + "prompt_embeds": None, + "prompt_embeds_mask": None, + "negative_prompt_embeds": None, + "negative_prompt_embeds_mask": None, }, ) prompt_embeds = prompt_fields["prompt_embeds"] @@ -1040,9 +1016,16 @@ class QwenImagePipeline( prompt = None if negative_prompt_embeds is not None: negative_prompt = None - true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale + true_cfg_scale = common_sampling_params.true_cfg_scale or 4.0 if common_sampling_params.guidance_scale_provided: guidance_scale = common_sampling_params.guidance_scale + else: + guidance_scale = 1.0 + + latents = req.collate_request_tensors("latents", None) + output_type = common_sampling_params.output_type or "pil" + attention_kwargs = None + callback_on_step_end_tensor_inputs = ["latents"] ctx = self._prepare_generation_context( prompt=prompt, @@ -1087,6 +1070,7 @@ class QwenImagePipeline( ) self._current_timestep = None + result = self._decode_latents(latents, height, width, output_type) return split_diffusion_output_by_request( result, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index e594d5ad8..3d6e937c7 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -7,7 +7,7 @@ import logging import math import os from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -677,30 +677,7 @@ class QwenImageEditPipeline( def interrupt(self): return self._interrupt - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - image: PIL.Image.Image | torch.Tensor | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_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, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 1024, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Forward pass for image editing.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. @@ -731,32 +708,28 @@ class QwenImageEditPipeline( height = req.sampling_params.height width = req.sampling_params.width else: - # fallback to run pre-processing in pipeline (debug only) - image_size = image[0].size if isinstance(image, list) else image.size - calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1]) - height = height or calculated_height - width = width or calculated_width + raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.") - height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) - - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): - image = self.image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image - image = self.image_processor.preprocess(image, calculated_height, calculated_width) - image = image.unsqueeze(2) - - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + max_sequence_length = req.sampling_params.max_sequence_length or 1024 + generator = req.sampling_params.generator + true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0 if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 1.0 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + prompt_embeds_mask = None + negative_prompt_embeds_mask = None + output_type = req.sampling_params.output_type or "pil" + attention_kwargs = None + callback_on_step_end_tensor_inputs = ["latents"] # 1. check inputs # 2. encode prompts @@ -783,16 +756,9 @@ class QwenImageEditPipeline( self._current_timestep = None 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] + batch_size = 1 - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) + has_neg_prompt = negative_prompt is not None do_true_cfg = true_cfg_scale > 1 and has_neg_prompt self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index bf63805b3..a11637163 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -5,7 +5,7 @@ import json import logging import os from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -623,30 +623,7 @@ class QwenImageEditPlusPipeline( def interrupt(self): return self._interrupt - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_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, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 1024, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Forward pass for image editing with support for multiple images.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. @@ -674,55 +651,35 @@ class QwenImageEditPlusPipeline( ): condition_images = additional_information.get("condition_images") vae_images = additional_information.get("vae_images") - condition_image_sizes = additional_information.get("condition_image_sizes") vae_image_sizes = additional_information.get("vae_image_sizes") calculated_height = additional_information.get("calculated_height") calculated_width = additional_information.get("calculated_width") - height = req.sampling_params.height - width = req.sampling_params.width + height = req.sampling_params.height or calculated_height + width = req.sampling_params.width or calculated_width + image = None else: - # fallback to run pre-processing in pipeline (debug only) - if image is None: - raise ValueError("Image is required for QwenImageEditPlusPipeline") + raise RuntimeError("Missing preprocess images that should have been created by the preprocess function.") - if not isinstance(image, list): - image = [image] - - image_size = image[0].size - calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1]) - height = height or calculated_height - width = width or calculated_width - - height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) - - condition_images = [] - vae_images = [] - condition_image_sizes = [] - vae_image_sizes = [] - - for img in image: - image_width, image_height = img.size - condition_width, condition_height = calculate_dimensions( - CONDITION_IMAGE_SIZE, image_width / image_height - ) - vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height) - condition_image_sizes.append((condition_width, condition_height)) - vae_image_sizes.append((vae_width, vae_height)) - condition_images.append(self.image_processor.resize(img, condition_height, condition_width)) - vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2)) - - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + max_sequence_length = req.sampling_params.max_sequence_length or 1024 + generator = req.sampling_params.generator + true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0 if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 1.0 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + prompt_embeds_mask = None + negative_prompt_embeds_mask = None + output_type = req.sampling_params.output_type or "pil" + attention_kwargs = None + callback_on_step_end_tensor_inputs = ["latents"] # 1. check inputs # 2. encode prompts @@ -749,16 +706,9 @@ class QwenImageEditPlusPipeline( self._current_timestep = None 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] + batch_size = 1 - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) + has_neg_prompt = negative_prompt is not None do_true_cfg = true_cfg_scale > 1 and has_neg_prompt self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index f7a9fc4b1..e2fe67976 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -7,7 +7,7 @@ import logging import math import os from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -648,31 +648,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as def interrupt(self): return self._interrupt - def forward( - self, - req: DiffusionRequestBatch, - image: PIL.Image.Image | torch.Tensor | None = None, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - true_cfg_scale: float = 4.0, - layers: int | None = 4, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float | None = None, - 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, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int = 1024, - resolution: int = 640, - cfg_normalize: bool = False, - use_en_prompt: bool = False, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Forward pass for image layered.""" # 1. Get preprocessed image from request (pre-processing is done in DiffusionEngine) @@ -688,26 +664,27 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") - layers = req.sampling_params.layers if req.sampling_params.layers is not None else layers - resolution = req.sampling_params.resolution if req.sampling_params.resolution is not None else resolution - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - cfg_normalize = ( - req.sampling_params.cfg_normalize if req.sampling_params.cfg_normalize is not None else cfg_normalize - ) - use_en_prompt = ( - req.sampling_params.use_en_prompt if req.sampling_params.use_en_prompt is not None else use_en_prompt - ) - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale + layers = req.sampling_params.layers + max_sequence_length = req.sampling_params.max_sequence_length or 1024 + cfg_normalize = req.sampling_params.cfg_normalize + use_en_prompt = req.sampling_params.use_en_prompt + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + generator = req.sampling_params.generator + true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0 + guidance_scale = None if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + prompt_embeds_mask = None + negative_prompt_embeds_mask = None + output_type = req.sampling_params.output_type or "pil" + attention_kwargs = None if not isinstance(first_prompt, str) and "preprocessed_image" in ( additional_information := first_prompt.get("additional_information", {}) @@ -720,25 +697,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as height = req.sampling_params.height width = req.sampling_params.width else: - # fallback to run pre-processing in pipeline (debug only) - if isinstance(image, PIL.Image.Image) and image.mode != "RGBA": - image = image.convert("RGBA") - image_size = image[0].size if isinstance(image, list) else image.size - assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}" - calculated_width, calculated_height = calculate_dimensions( - resolution * resolution, image_size[0] / image_size[1] - ) - height = calculated_height - width = calculated_width - - height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) - - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): - image = self.image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image - image = self.image_processor.preprocess(image, calculated_height, calculated_width) - image = image.unsqueeze(2) - image = image.to(dtype=self.text_encoder.dtype) + raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.") # 2. check inputs self.check_inputs( @@ -761,16 +720,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as # 3. encode prompot & negative prompt if prompt is None or prompt == "" or prompt == " ": prompt = self.get_image_caption(prompt_image, use_en_prompt=use_en_prompt, device=self.device) - 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] + batch_size = 1 - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) + has_neg_prompt = negative_prompt is not None if true_cfg_scale > 1 and not has_neg_prompt: logger.warning( diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py index e45614a42..b2852e524 100644 --- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py +++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py @@ -628,57 +628,39 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro return latents - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] = "", - prompt_2: str | list[str] = "", - prompt_3: str | list[str] = "", - negative_prompt: str | list[str] = "", - negative_prompt_2: str | list[str] = "", - negative_prompt_3: str | list[str] = "", - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 28, - sigmas: list[float] | None = None, - 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, - negative_prompt_embeds: torch.Tensor | None = None, - pooled_prompt_embeds: torch.Tensor | None = None, - negative_pooled_prompt_embeds: torch.Tensor | None = None, - max_sequence_length: int = 256, - ) -> list[DiffusionOutput]: + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. sampling_params_list = req.sampling_params_list common_sampling_params = sampling_params_list[0] - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt - negative_prompt = [ - "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts - ] or negative_prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] + prompt_2 = "" + prompt_3 = "" + negative_prompt_2 = "" + negative_prompt_3 = "" + prompt_embeds = None + negative_prompt_embeds = None height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor - sigmas = common_sampling_params.sigmas or sigmas - max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length - num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps + sigmas = common_sampling_params.sigmas + max_sequence_length = common_sampling_params.max_sequence_length or 256 + num_inference_steps = common_sampling_params.num_inference_steps or 28 num_images_per_prompt = ( - common_sampling_params.num_outputs_per_prompt - if common_sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1 ) - if generator is None: - generator = req.collate_request_generators(num_images_per_prompt, generator) - latents = req.collate_request_tensors("latents", latents) + generator = req.collate_request_generators(num_images_per_prompt, None) + latents = req.collate_request_tensors("latents", None) + output_type = common_sampling_params.output_type or "pil" + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( req.prompts, { "prompt_embeds": prompt_embeds, "negative_prompt_embeds": negative_prompt_embeds, - "pooled_prompt_embeds": pooled_prompt_embeds, - "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds, + "pooled_prompt_embeds": None, + "negative_pooled_prompt_embeds": None, }, ) prompt_embeds = prompt_fields["prompt_embeds"] @@ -717,9 +699,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro self._current_timestep = None self._interrupt = False - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): + if prompt is not None: batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] @@ -771,7 +751,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro ) self._current_timestep = None - if self.output_type == "latent": + if output_type == "latent": image = latents else: # Ensure the latents are the same dtype as the VAE for decode diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py index 23b61209a..41f12f654 100644 --- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py +++ b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py @@ -385,37 +385,21 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov latents = latents * self.scheduler.init_noise_sigma return latents - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - audio_end_in_s: float | None = None, - audio_start_in_s: float = 0.0, - guidance_scale: float = 7.0, - num_waveforms_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - output_type: str = "np", - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """ Generate audio from text prompt. Args: - req: OmniDiffusionRequest containing generation parameters - prompt: Text prompt for audio generation - negative_prompt: Negative prompt for CFG - audio_end_in_s: Audio end time in seconds (max ~47s for stable-audio-open-1.0) - audio_start_in_s: Audio start time in seconds - guidance_scale: CFG scale - num_waveforms_per_prompt: Number of audio outputs per prompt - generator: Random generator for reproducibility - latents: Pre-generated latents - prompt_embeds: Pre-computed prompt embeddings - negative_prompt_embeds: Pre-computed negative prompt embeddings - output_type: Output format ("np", "pt", or "latent") + req: OmniDiffusionRequest containing generation parameters. + The `req.sampling_params.extra_args` can include the following keys: + - audio_start_in_s (`float`, *optional*, defaults to 0.0): + Start time of the audio in seconds. + - audio_end_in_s (`float`, *optional*): + End time of the audio in seconds. + - num_waveforms_per_prompt (`int`, *optional*, defaults to 1): + Number of audio outputs per prompt. + - output_type (`str`, *optional*, defaults to "np"): + Output format ("np", "pt", or "latent"). Returns: DiffusionOutput containing generated audio @@ -423,28 +407,33 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov # Extract from request # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - num_inference_steps = req.sampling_params.num_inference_steps - if num_inference_steps is None: - num_inference_steps = 100 # Default steps - num_inference_steps = req.sampling_params.num_inference_steps - if num_inference_steps is None: - num_inference_steps = 50 # Default steps + + prompt_embeds = None + negative_prompt_embeds = None + + num_inference_steps = req.sampling_params.num_inference_steps or 50 if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 7.0 - if generator is None: - generator = req.sampling_params.generator + generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) + latents = req.sampling_params.latents # Get audio duration from request extra params or defaults - audio_start_in_s = req.sampling_params.extra_args.get("audio_start_in_s", audio_start_in_s) - audio_end_in_s = req.sampling_params.extra_args.get("audio_end_in_s", audio_end_in_s) + audio_start_in_s: float = req.sampling_params.extra_args.get("audio_start_in_s", 0.0) + audio_end_in_s: float | None = req.sampling_params.extra_args.get("audio_end_in_s", None) + num_waveforms_per_prompt: int = req.sampling_params.extra_args.get( + "num_waveforms_per_prompt", req.sampling_params.num_outputs_per_prompt or 1 + ) + output_type = req.sampling_params.output_type or "np" # Calculate audio length downsample_ratio = self.vae.hop_length @@ -474,12 +463,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov ) # Determine batch size - 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] + batch_size = len(prompt) device = self.device do_classifier_free_guidance = guidance_scale > 1.0 @@ -500,7 +484,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov audio_start_in_s, audio_end_in_s, device, - do_classifier_free_guidance and (negative_prompt is not None or negative_prompt_embeds is not None), + do_classifier_free_guidance and negative_prompt is not None, batch_size, ) @@ -515,7 +499,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov ) # Handle CFG without negative prompt - if do_classifier_free_guidance and negative_prompt_embeds is None and negative_prompt is None: + if do_classifier_free_guidance and negative_prompt is None: negative_text_audio_duration_embeds = torch.zeros_like(text_audio_duration_embeds) text_audio_duration_embeds = torch.cat( [negative_text_audio_duration_embeds, text_audio_duration_embeds], diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index b2870549c..5d7c8a281 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -540,38 +540,27 @@ class Wan22Pipeline( return latents - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | None = None, - negative_prompt: str | None = None, - height: int = 480, - width: int = 832, - num_inference_steps: int = 40, - guidance_scale: float | tuple[float, float] = 4.0, - frame_num: int = 81, - output_type: str | None = "np", - generator: torch.Generator | list[torch.Generator] | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - attention_kwargs: dict | None = None, - **kwargs, - ) -> DiffusionOutput: - # Get parameters from request or arguments + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + prompt: str | None = None + negative_prompt: str | None = None + prompt_embeds: torch.Tensor | None = None + negative_prompt_embeds: torch.Tensor | None = None if len(req.prompts) > 1: raise ValueError( """This model only supports a single prompt, not a batched request.""", """Please pass in a single prompt object or string, or a single-item list.""", ) - if len(req.prompts) == 1: # If req.prompt is empty, default to prompt & neg_prompt in param list - prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") - negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") - if prompt is None and prompt_embeds is None: - raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.") + if len(req.prompts) == 1: + first_prompt = req.prompts[0] + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num + if not prompt: + raise ValueError("Prompt is required for Wan2.2 generation.") + + height = req.sampling_params.height or 480 + width = req.sampling_params.width or 832 + num_frames = req.sampling_params.num_frames or 81 # Ensure dimensions are compatible with VAE and patch size # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size) @@ -579,11 +568,16 @@ class Wan22Pipeline( mod_value = self.vae_scale_factor_spatial * patch_size[1] # 16*2=32 for TI2V, 8*2=16 for I2V height = (height // mod_value) * mod_value width = (width // mod_value) * mod_value - num_steps = req.sampling_params.num_inference_steps or num_inference_steps + num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps # Respect per-request guidance_scale when explicitly provided. if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 4.0 + + output_type = req.sampling_params.output_type or "np" + attention_kwargs: dict | None = None guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0] guidance_high = ( @@ -634,8 +628,7 @@ class Wan22Pipeline( dtype = self.text_encoder.dtype # Seed / generator - if generator is None: - generator = req.sampling_params.generator + generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) @@ -644,24 +637,16 @@ class Wan22Pipeline( current_omni_platform.synchronize() _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - if prompt_embeds is None: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) - else: - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - if negative_prompt_embeds is not None: - negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) - elif guidance_low > 1.0 or guidance_high > 1.0: - raise ValueError( - "negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1." - ) + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, + num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, + max_sequence_length=req.sampling_params.max_sequence_length or 512, + device=device, + dtype=dtype, + ) + if DEBUG_PERF: current_omni_platform.synchronize() _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index fc2d79cec..65f74b7f9 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -429,66 +429,59 @@ class Wan22I2VPipeline( quant_config = getattr(self.od_config, "quantization_config", None) return create_transformer_from_config(config, quant_config=quant_config) - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | None = None, - negative_prompt: str | None = None, - image: PIL.Image.Image | torch.Tensor | None = None, - height: int = 480, - width: int = 832, - num_inference_steps: int = 40, - guidance_scale: float | tuple[float, float] = 5.0, - frame_num: int = 81, - output_type: str | None = "np", - generator: torch.Generator | list[torch.Generator] | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - image_embeds: torch.Tensor | None = None, - last_image: PIL.Image.Image | torch.Tensor | None = None, - attention_kwargs: dict | None = None, - **kwargs, - ) -> DiffusionOutput: - # Get parameters from request or arguments + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + prompt: str | None = None + negative_prompt: str | None = None + prompt_embeds: torch.Tensor | None = None + negative_prompt_embeds: torch.Tensor | None = None + image_embeds: torch.Tensor | None = None if len(req.prompts) > 1: raise ValueError( """This model only supports a single prompt, not a batched request.""", """Please pass in a single prompt object or string, or a single-item list.""", ) - if len(req.prompts) == 1: # If req.prompt is empty, default to prompt & neg_prompt in param list - prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") - negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") - if prompt is None and prompt_embeds is None: - raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.") + if len(req.prompts) == 1: + first_prompt = req.prompts[0] + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") + + if not prompt: + raise ValueError("Prompt is required for Wan2.2 generation.") # Get image from request - if image is None: - multi_modal_data = ( - req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None - ) - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if raw_image is None: - raise ValueError("Image is required for I2V generation.") - if isinstance(raw_image, list): - if len(raw_image) > 1: - logger.warning( - """Received a list of image. Only a single image is supported by this model.""" - """Taking only the first image for now.""" - ) - raw_image = raw_image[0] - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image) - else: - image = cast(PIL.Image.Image | torch.Tensor, raw_image) + multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if raw_image is None: + raise ValueError("Image is required for I2V generation.") + if isinstance(raw_image, list): + if len(raw_image) > 1: + logger.warning( + """Received a list of image. Only a single image is supported by this model.""" + """Taking only the first image for now.""" + ) + raw_image = raw_image[0] + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image) + else: + image = cast(PIL.Image.Image | torch.Tensor, raw_image) - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_frames = req.sampling_params.num_frames or frame_num - num_steps = req.sampling_params.num_inference_steps or num_inference_steps + last_image: PIL.Image.Image | torch.Tensor | None = None + if multi_modal_data is not None: + last_image = multi_modal_data.get("last_image", None) + + height = req.sampling_params.height or 480 + width = req.sampling_params.width or 832 + num_frames = req.sampling_params.num_frames or 81 + num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps # Respect per-request guidance_scale when explicitly provided. if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 5.0 + + output_type = req.sampling_params.output_type or "np" + attention_kwargs: dict | None = None # Handle guidance scales guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0] @@ -533,8 +526,7 @@ class Wan22I2VPipeline( dtype = self.transformer.dtype # Generator setup - if generator is None: - generator = req.sampling_params.generator + generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) @@ -544,20 +536,15 @@ class Wan22I2VPipeline( _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - if prompt_embeds is None: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) - else: - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - if negative_prompt_embeds is not None: - negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, + num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, + max_sequence_length=req.sampling_params.max_sequence_length or 512, + device=device, + dtype=dtype, + ) if DEBUG_PERF: current_omni_platform.synchronize() diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 104576f41..b77d2265c 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -24,7 +24,7 @@ from typing import Any, ClassVar import PIL.Image import torch import torch.nn as nn -from diffusers.image_processor import PipelineImageInput, VaeImageProcessor +from diffusers.image_processor import VaeImageProcessor from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils import logging from diffusers.utils.torch_utils import randn_tensor @@ -410,124 +410,21 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen def interrupt(self): return self._interrupt - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - image: PipelineImageInput = None, - strength: float = 0.6, - height: int = 1024, - width: int = 1024, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 5.0, - cfg_normalization: bool = False, - cfg_truncation: float = 1.0, - negative_prompt: str | list[str] | None = None, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: list[torch.FloatTensor] | None = None, - negative_prompt_embeds: list[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 = 512, - ) -> 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. - image (`PipelineImageInput`, *optional*): - The image to use for img2img generation. If provided, the pipeline - will perform img2img instead of text-to-image. - strength (`float`, *optional*, defaults to 0.6): - Indicates extent to transform the reference `image`. Must be between 0 and 1. - height (`int`, *optional*, defaults to 1024): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 1024): - The width in pixels of the generated image. - 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. - guidance_scale (`float`, *optional*, defaults to 5.0): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 0`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - cfg_normalization (`bool`, *optional*, defaults to False): - Whether to apply configuration normalization. - cfg_truncation (`float`, *optional*, defaults to 1.0): - The truncation value for configuration. - 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 - less than or equal to `0`). - 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 (`list[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 (`list[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.stable_diffusion.ZImagePipelineOutput`] 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`, *optional*, defaults to 512): - Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the - generated images. - """ + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - # Handle img2img: extract image from request - if image is None and req.prompts: + prompt_embeds = None + negative_prompt_embeds = None + + image = None + if req.prompts: if len(req.prompts) > 1: logger.warning( "This model only supports a single prompt for img2img, not a batched request. " @@ -542,10 +439,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen else: image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image - # strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter explicit_strength = req.sampling_params.strength is not None - if explicit_strength: - strength = req.sampling_params.strength + strength = req.sampling_params.strength if explicit_strength else 0.6 if explicit_strength and image is None: logger.warning( "strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. " @@ -556,20 +451,24 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen if image is not None and strength is not None and (strength < 0 or strength > 1): raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}") - 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 + height = req.sampling_params.height or 1024 + width = req.sampling_params.width or 1024 + num_inference_steps = req.sampling_params.num_inference_steps or 50 generator = req.sampling_params.generator - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_rescale is not None else guidance_scale - ) + sigmas = req.sampling_params.sigmas + max_sequence_length = req.sampling_params.max_sequence_length or 512 + guidance_scale = req.sampling_params.guidance_scale num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + + cfg_normalization = req.sampling_params.cfg_normalize + cfg_truncation = req.sampling_params.extra_args.get("cfg_truncation", 1.0) + 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 = ["latents"] + output_type = req.sampling_params.output_type or "pil" vae_scale = self.vae_scale_factor * 2 if height % vae_scale != 0: @@ -591,33 +490,20 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen self._cfg_normalization = cfg_normalization self._cfg_truncation = cfg_truncation # 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 = len(prompt_embeds) + batch_size = len(prompt) - # If prompt_embeds is provided and prompt is None, skip encoding - if prompt_embeds is not None and prompt is None: - if self.do_classifier_free_guidance and negative_prompt_embeds is None: - raise ValueError( - "When `prompt_embeds` is provided without `prompt`, " - "`negative_prompt_embeds` must also be provided for classifier-free guidance." - ) - else: - ( - prompt_embeds, - negative_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - device=device, - max_sequence_length=max_sequence_length, - ) + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + ) # 4. Prepare latent variables num_channels_latents = self.transformer.in_channels