diff --git a/src/diffusers/models/transformers/transformer_ovis_image.py b/src/diffusers/models/transformers/transformer_ovis_image.py index 7a9df427e0b9..44723bc44fd0 100644 --- a/src/diffusers/models/transformers/transformer_ovis_image.py +++ b/src/diffusers/models/transformers/transformer_ovis_image.py @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import logging +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_adjust_dtype_for_device, maybe_allow_in_graph from ..attention import AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn @@ -473,6 +473,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -480,6 +481,7 @@ def forward( timestep: torch.LongTensor = None, img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, ) -> torch.Tensor | Transformer2DModelOutput: """ @@ -496,6 +498,10 @@ def forward( The position ids for image tokens. txt_ids (`torch.Tensor`): The position ids for text tokens. + 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). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain tuple. @@ -538,6 +544,7 @@ def forward( encoder_hidden_states, temb, image_rotary_emb, + joint_attention_kwargs, ) else: @@ -546,6 +553,7 @@ def forward( encoder_hidden_states=encoder_hidden_states, temb=temb, image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, ) for index_block, block in enumerate(self.single_transformer_blocks): @@ -556,6 +564,7 @@ def forward( encoder_hidden_states, temb, image_rotary_emb, + joint_attention_kwargs, ) else: @@ -564,6 +573,7 @@ def forward( encoder_hidden_states=encoder_hidden_states, temb=temb, image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, ) hidden_states = self.norm_out(hidden_states, temb) diff --git a/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py b/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py index c8ff8227f27e..b22f2f0cec2d 100644 --- a/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py +++ b/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py @@ -176,7 +176,6 @@ def __init__( self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) self.system_prompt = "Describe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background: " self.user_prompt_begin_id = 28 - self.tokenizer_max_length = 256 + self.user_prompt_begin_id self.default_sample_size = 128 def _get_messages( @@ -202,11 +201,12 @@ def _get_ovis_prompt_embeds( self, prompt: str | list[str] = None, num_images_per_prompt: int = 1, + max_sequence_length: int = 256, device: torch.device | None = None, dtype: torch.dtype | None = None, ): device = device or self._execution_device - dtype = dtype or self.text_encoder.dtype + dtype = dtype or (self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype) messages = self._get_messages(prompt) batch_size = len(messages) @@ -215,7 +215,7 @@ def _get_ovis_prompt_embeds( messages, padding="max_length", truncation=True, - max_length=self.tokenizer_max_length, + max_length=max_sequence_length + self.user_prompt_begin_id, return_tensors="pt", add_special_tokens=False, ) @@ -237,41 +237,86 @@ def _get_ovis_prompt_embeds( return prompt_embeds + def _prepare_prompt_embeds( + self, + prompt: str | list[str], + prompt_embeds: torch.FloatTensor | None, + num_images_per_prompt: int, + max_sequence_length: int, + device: torch.device, + ): + if prompt_embeds is None: + prompt_embeds = self._get_ovis_prompt_embeds( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + else: + dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + batch_size, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype + text_ids = torch.zeros(prompt_embeds.shape[1], 3) + text_ids[..., 1] = text_ids[..., 1] + torch.arange(prompt_embeds.shape[1])[None, :] + text_ids[..., 2] = text_ids[..., 2] + torch.arange(prompt_embeds.shape[1])[None, :] + text_ids = text_ids.to(device=device, dtype=dtype) + return prompt_embeds, text_ids + def encode_prompt( self, prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, device: torch.device | None = None, num_images_per_prompt: int = 1, + max_sequence_length: int = 256, prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, ): r""" Args: - prompt (`str`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. Used only when `do_classifier_free_guidance` + is `True`. If not defined, an empty string is used. + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to also encode the `negative_prompt` for classifier-free guidance. device: (`torch.device`): torch device num_images_per_prompt (`int`): number of images that should be generated per prompt + max_sequence_length (`int`, *optional*, defaults to 256): + Maximum sequence length to use for the `prompt`. prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. If not provided, they are generated from `negative_prompt`. """ device = device or self._execution_device - if prompt_embeds is None: - prompt_embeds = self._get_ovis_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, + prompt_embeds, text_ids = self._prepare_prompt_embeds( + prompt, prompt_embeds, num_images_per_prompt, max_sequence_length, device + ) + + negative_text_ids = None + if do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = "" + if isinstance(negative_prompt, str): + batch_size = prompt_embeds.shape[0] // num_images_per_prompt + negative_prompt = batch_size * [negative_prompt] + negative_prompt_embeds, negative_text_ids = self._prepare_prompt_embeds( + negative_prompt, negative_prompt_embeds, num_images_per_prompt, max_sequence_length, device ) - dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype - text_ids = torch.zeros(prompt_embeds.shape[1], 3) - text_ids[..., 1] = text_ids[..., 1] + torch.arange(prompt_embeds.shape[1])[None, :] - text_ids[..., 2] = text_ids[..., 2] + torch.arange(prompt_embeds.shape[1])[None, :] - text_ids = text_ids.to(device=device, dtype=dtype) - return prompt_embeds, text_ids + return prompt_embeds, negative_prompt_embeds, text_ids, negative_text_ids def check_inputs( self, @@ -516,6 +561,7 @@ def __call__( max_sequence_length=max_sequence_length, ) + self._guidance_scale = guidance_scale self._joint_attention_kwargs = joint_attention_kwargs self._current_timestep = None self._interrupt = False @@ -533,23 +579,19 @@ def __call__( do_classifier_free_guidance = guidance_scale > 1 ( prompt_embeds, + negative_prompt_embeds, text_ids, + negative_text_ids, ) = self.encode_prompt( prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, device=device, num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, ) - if do_classifier_free_guidance: - ( - negative_prompt_embeds, - negative_text_ids, - ) = self.encode_prompt( - prompt=negative_prompt, - prompt_embeds=negative_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) # 4. Prepare latent variables num_channels_latents = self.transformer.config.in_channels // 4 @@ -609,6 +651,7 @@ def __call__( encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] @@ -620,6 +663,7 @@ def __call__( encoder_hidden_states=negative_prompt_embeds, txt_ids=negative_text_ids, img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] noise_pred = neg_noise_pred + guidance_scale * (noise_pred - neg_noise_pred) diff --git a/tests/pipelines/ovis_image/test_ovis_image.py b/tests/pipelines/ovis_image/test_ovis_image.py new file mode 100644 index 000000000000..ac1670535363 --- /dev/null +++ b/tests/pipelines/ovis_image/test_ovis_image.py @@ -0,0 +1,153 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch +from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model + +from diffusers import ( + AutoencoderKL, + FlowMatchEulerDiscreteScheduler, + OvisImagePipeline, + OvisImageTransformer2DModel, +) + +from ...testing_utils import torch_device +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +class OvisImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = OvisImagePipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + supports_dduf = False + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = OvisImageTransformer2DModel( + patch_size=1, + in_channels=4, + out_channels=4, + num_layers=1, + num_single_layers=1, + attention_head_dim=16, + num_attention_heads=2, + joint_attention_dim=32, + axes_dims_rope=(4, 4, 8), + ) + torch.manual_seed(0) + vae = AutoencoderKL( + sample_size=32, + in_channels=3, + out_channels=3, + block_out_channels=(4,), + layers_per_block=1, + latent_channels=1, + norm_num_groups=1, + use_quant_conv=False, + use_post_quant_conv=False, + shift_factor=0.0609, + scaling_factor=1.5035, + ) + scheduler = FlowMatchEulerDiscreteScheduler() + tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + torch.manual_seed(0) + text_encoder = Qwen3Model( + Qwen3Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + vocab_size=tokenizer.vocab_size + 4, + max_position_embeddings=512, + ) + ) + return { + "scheduler": scheduler, + "vae": vae, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "transformer": transformer, + } + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + return { + "prompt": "a cat", + "negative_prompt": "bad quality", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 2.0, + "height": 16, + "width": 16, + "max_sequence_length": 32, + "output_type": "np", + } + + def test_inference(self): + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + generated_image = image[0] + self.assertEqual(generated_image.shape, (16, 16, 3)) + self.assertTrue(np.isfinite(image).all()) + + def test_guidance_scale_is_set(self): + # The `guidance_scale` property reads `self._guidance_scale`, which `__call__` must initialize. + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + pipe(**inputs) + assert pipe.guidance_scale == inputs["guidance_scale"] + + def test_max_sequence_length_is_used(self): + # `max_sequence_length` should bound the encoded prompt length. + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + embeds_16 = pipe.encode_prompt( + "a cat", do_classifier_free_guidance=False, device=torch_device, max_sequence_length=16 + )[0] + embeds_32 = pipe.encode_prompt( + "a cat", do_classifier_free_guidance=False, device=torch_device, max_sequence_length=32 + )[0] + assert embeds_16.shape[1] == 16 + assert embeds_32.shape[1] == 32