ovis_image model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules. AGENTS.md is referenced by .ai/review-rules.md but is not present in this checkout; the other referenced rule files were read and applied.

Files/categories reviewed: target pipeline/model files, public imports and lazy loading, config/serialization, dtype/device/offload paths, attention processor behavior, docs, and fast/slow test coverage.

Duplicate search status: gh search hit the GitHub API rate limit, so I checked GitHub web issue/PR searches for OvisImage, ovis_image, OvisImageTransformer2DModel AttentionMixin, and OvisImagePipeline num_images_per_prompt. I did not find an exact Ovis duplicate. Related but not duplicate: #12186 covers the same missing-AttentionMixin pattern for WanVACETransformer3DModel.

Issue 1: Transformer does not expose attention processor APIs

Affected code:

Problem:
OvisImageTransformer2DModel defines OvisImageAttention modules but does not inherit AttentionMixin. That leaves the model without the standard attn_processors, set_attn_processor, fuse_qkv_projections, and unfuse_qkv_projections APIs expected by related transformer families.

Impact:
Users and tests cannot swap attention processors, inspect processors, or use QKV fusion through the model-level API.

Reproduction:

from diffusers import OvisImageTransformer2DModel

print(hasattr(OvisImageTransformer2DModel, "set_attn_processor"))
print(hasattr(OvisImageTransformer2DModel, "fuse_qkv_projections"))
assert hasattr(OvisImageTransformer2DModel, "set_attn_processor")

Relevant precedent:

Suggested fix:

from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward

class OvisImageTransformer2DModel(
    ModelMixin,
    ConfigMixin,
    PeftAdapterMixin,
    FromOriginalModelMixin,
    CacheMixin,
    AttentionMixin,
):
    ...

Issue 2: joint_attention_kwargs is accepted but never reaches attention

Affected code:

Problem:
The pipeline exposes joint_attention_kwargs, but OvisImageTransformer2DModel.forward() does not accept it, and the pipeline transformer calls do not pass it. The block classes already accept joint_attention_kwargs, so the plumbing is incomplete.

Impact:
Any user-provided attention kwargs are silently ignored by the pipeline. Direct model calls with the same argument fail.

Reproduction:

import inspect
from diffusers import OvisImagePipeline, OvisImageTransformer2DModel

print("pipeline:", "joint_attention_kwargs" in inspect.signature(OvisImagePipeline.__call__).parameters)
print("model:", "joint_attention_kwargs" in inspect.signature(OvisImageTransformer2DModel.forward).parameters)
assert "joint_attention_kwargs" in inspect.signature(OvisImageTransformer2DModel.forward).parameters

Relevant precedent:

Suggested fix:

def forward(..., joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True):
    ...
    encoder_hidden_states, hidden_states = block(
        hidden_states=hidden_states,
        encoder_hidden_states=encoder_hidden_states,
        temb=temb,
        image_rotary_emb=image_rotary_emb,
        joint_attention_kwargs=joint_attention_kwargs,
    )

Also pass joint_attention_kwargs=self.joint_attention_kwargs in both pipeline transformer calls.

Issue 3: guidance_scale property is never initialized

Affected code:

Problem:
OvisImagePipeline.guidance_scale returns self._guidance_scale, but __call__ never assigns self._guidance_scale = guidance_scale.

Impact:
Callbacks or downstream code that read pipe.guidance_scale during generation can hit an AttributeError or stale state, unlike related pipelines.

Reproduction:

import inspect
from diffusers import OvisImagePipeline

source = inspect.getsource(OvisImagePipeline.__call__)
print("self._guidance_scale = guidance_scale" in source)
assert "self._guidance_scale = guidance_scale" in source

Relevant precedent:

Suggested fix:

self._guidance_scale = guidance_scale
self._joint_attention_kwargs = joint_attention_kwargs
self._current_timestep = None
self._interrupt = False

Issue 4: Batched prompts break with default negative prompt under CFG

Affected code:

Problem:
With prompt as a list and default negative_prompt="", positive embeddings are batched to len(prompt), but negative embeddings are encoded as batch size 1. CFG then calls the transformer with mismatched latent and negative prompt batch sizes.

Impact:
The default CFG path fails for normal batched text-to-image usage unless users manually pass a negative prompt list of matching length.

Reproduction:

import torch
from diffusers import OvisImagePipeline

pipe = OvisImagePipeline.__new__(OvisImagePipeline)
pipe.text_encoder = type("E", (), {"dtype": torch.float32})()
pipe.transformer = type("T", (), {"dtype": torch.float32})()

def fake_get_ovis_prompt_embeds(prompt, num_images_per_prompt=1, device=None, dtype=None):
    prompt = [prompt] if isinstance(prompt, str) else prompt
    return torch.zeros(len(prompt) * num_images_per_prompt, 4, 8)

pipe._get_ovis_prompt_embeds = fake_get_ovis_prompt_embeds
pos, _ = pipe.encode_prompt(["cat", "dog"], device=torch.device("cpu"))
neg, _ = pipe.encode_prompt("", device=torch.device("cpu"))
print(pos.shape[0], neg.shape[0])
assert pos.shape[0] == neg.shape[0]

Relevant precedent:

Suggested fix:

if do_classifier_free_guidance and negative_prompt_embeds is None:
    if negative_prompt is None:
        negative_prompt = ""
    if isinstance(negative_prompt, str):
        negative_prompt = [negative_prompt] * batch_size
    elif len(negative_prompt) != batch_size:
        raise ValueError(
            f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has batch size {batch_size}."
        )

Issue 5: Precomputed prompt_embeds are not moved or repeated

Affected code:

Problem:
When prompt_embeds is supplied, encode_prompt() does not move it to the execution device/dtype and does not repeat it for num_images_per_prompt. The pipeline still prepares latents for batch_size * num_images_per_prompt.

Impact:
Precomputed embeddings can fail with device mismatches on GPU/offload paths and batch mismatches when generating multiple images per prompt.

Reproduction:

import torch
from diffusers import OvisImagePipeline

pipe = OvisImagePipeline.__new__(OvisImagePipeline)
pipe.text_encoder = None
pipe.transformer = type("T", (), {"dtype": torch.float16})()

embeds = torch.randn(1, 4, 8, dtype=torch.float32)
out, ids = pipe.encode_prompt(None, device=torch.device("meta"), prompt_embeds=embeds, num_images_per_prompt=2)

print(out.shape, out.device, ids.device)
assert out.shape[0] == 2
assert out.device == ids.device

Relevant precedent:

Suggested fix:

batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]

if prompt_embeds is None:
    prompt_embeds = self._get_ovis_prompt_embeds(...)

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)
_, 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)

Issue 6: max_sequence_length is validated but ignored

Affected code:

Problem:
__call__ accepts and validates max_sequence_length, but encode_prompt() has no such parameter and _get_ovis_prompt_embeds() always tokenizes with self.tokenizer_max_length.

Impact:
Users cannot reduce prompt sequence length for speed/memory, and the public argument is misleading.

Reproduction:

import inspect
from diffusers import OvisImagePipeline

assert "max_sequence_length" in inspect.signature(OvisImagePipeline.__call__).parameters
assert "max_sequence_length" in inspect.signature(OvisImagePipeline.encode_prompt).parameters

Relevant precedent:

Suggested fix:

def _get_ovis_prompt_embeds(..., max_sequence_length: int = 256):
    max_length = max_sequence_length + self.user_prompt_begin_id
    tokens = self.tokenizer(..., max_length=max_length, ...)
    prompt_embeds = prompt_embeds[:, self.user_prompt_begin_id : self.user_prompt_begin_id + max_sequence_length, :]

def encode_prompt(..., max_sequence_length: int = 256):
    ...

Then pass max_sequence_length=max_sequence_length from __call__.

Issue 7: No fast or slow tests cover ovis_image

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/ovis_image/__init__.py#L1

Problem:
tests/pipelines/ovis_image/ only contains an empty __init__.py, and there is no model test for OvisImageTransformer2DModel. No fast or slow tests reference OvisImage or ovis_image.

Impact:
The import surface, prompt batching, callback properties, attention APIs, serialization, and slow checkpoint path can regress without CI coverage.

Reproduction:

from pathlib import Path

paths = []
for path in Path("tests").rglob("test*.py"):
    text = path.read_text(encoding="utf-8", errors="ignore")
    if "OvisImage" in text or "ovis_image" in text or "ovis-image" in text.lower():
        paths.append(str(path))

print(paths)
assert paths, "No fast or slow Ovis tests found"

Relevant precedent:

Suggested fix:
Add a model fast test using ModelTesterMixin and AttentionTesterMixin, add a pipeline fast test with tiny synthetic components, and add at least one @slow pipeline smoke test for the published Ovis checkpoint. These tests should cover top-level imports, save/load, attention processor APIs, batched prompts with CFG, prompt_embeds, callback properties, and max_sequence_length.