ovis_image: fix guidance_scale / max_sequence_length / batched CFG / precomputed embeds + add pipeline test (huggingface#13944)

* ovis_image: fix guidance_scale / max_sequence_length / batched-CFG / precomputed embeds + add pipeline test

Addresses items 3/4/5/6/7 of huggingface#13630.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* ovis_image: complete pipeline review (huggingface#13630)

- thread joint_attention_kwargs through the transformer forward + blocks
  (item 2) and pass it from the pipeline's transformer calls.
- encode_prompt now returns both positive and negative embeds (the z_image /
  PixArt convention) so precomputed embeds work end-to-end and the prompt is
  encoded in a single call.
- switch the pipeline test to the full PipelineTesterMixin.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* ovis_image: remove now-unused tokenizer_max_length

The pipeline computes max_length inline from max_sequence_length, so the attribute is dead. Addresses review feedback.

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Author: 3 people

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,13 +473,15 @@ def __init__(
 
         self.gradient_checkpointing = False
 
+    @apply_lora_scale("joint_attention_kwargs")
     def forward(
         self,
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor = None,
         timestep: torch.LongTensor = None,
         img_ids: torch.Tensor = None,
         txt_ids: torch.Tensor = None,
+        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)

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)