kandinsky2_2 model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Duplicate search checked kandinsky2_2, KandinskyV22, affected class/file names, and the specific failure modes. Existing related items: #4183 covers the ControlNet guidance_scale <= 1 subset of Issue 2; #4818 is related to stale PriorEmb2Emb latents plumbing, but I did not find an exact duplicate for the current interpolate() failure.

Issue 1: PriorEmb2Emb.interpolate() cannot handle documented text entries

Affected code:

Problem:
interpolate() advertises list[str | PIL.Image.Image | torch.Tensor], but the string branch calls self(..., latents=latents) while KandinskyV22PriorEmb2EmbPipeline.__call__ has no latents parameter and requires image. The same method also returns torch.randn_like(image_emb) for negative_image_embeds, making the negative conditioning random and unrelated to negative_prompt.

Impact:
The documented mixed text/image interpolation workflow fails immediately for text entries. Image-only interpolation is nondeterministic on the negative branch and can change decoder CFG behavior across calls.

Reproduction:

from diffusers import KandinskyV22PriorEmb2EmbPipeline

pipe = KandinskyV22PriorEmb2EmbPipeline(
    prior=None, image_encoder=None, text_encoder=None,
    tokenizer=None, scheduler=None, image_processor=None,
)

try:
    pipe.interpolate(["a cat"], [1.0])
except TypeError as e:
    print(type(e).__name__, str(e).split("\n")[0])

Relevant precedent:
KandinskyV22PriorPipeline.interpolate() supports text entries by calling a compatible text-only prior path and derives the negative image embedding from a real negative/zero embedding:

Suggested fix:
Implement the text branch by reusing/copying the text-only prior logic from KandinskyV22PriorPipeline, or narrow the accepted input types and docs to image/tensor only. Replace torch.randn_like(image_emb) with a deterministic negative embedding, likely get_zero_embed(...) for the empty negative case and a real negative-prompt path when supported.

Issue 2: Decoder no-CFG paths do not repeat or cast conditioning tensors

Affected code:

Problem:
image_embeds, and hint for ControlNet, are repeated for num_images_per_prompt and moved to self.unet.dtype/device only inside the CFG branch. With guidance_scale <= 1, latents are sized for batch * num_images_per_prompt, but conditioning remains at the original batch size and dtype/device.

Impact:
No-CFG generation can fail for num_images_per_prompt > 1, half-precision pipelines, CPU-to-accelerator inputs, and ControlNet hints. The ControlNet subset is already reported in #4183.

Reproduction:

from types import SimpleNamespace
import torch
from diffusers import KandinskyV22Pipeline

class FakeUNet:
    dtype = torch.float32
    config = SimpleNamespace(in_channels=4)
    def __call__(self, sample, timestep, encoder_hidden_states=None, added_cond_kwargs=None, return_dict=False):
        assert added_cond_kwargs["image_embeds"].shape[0] == sample.shape[0], (
            added_cond_kwargs["image_embeds"].shape, sample.shape
        )
        return (torch.zeros(sample.shape[0], sample.shape[1] * 2, sample.shape[2], sample.shape[3]),)

class FakeScheduler:
    init_noise_sigma = 1.0
    config = SimpleNamespace(variance_type="learned")
    def set_timesteps(self, *args, **kwargs): self.timesteps = torch.tensor([1])
    def step(self, noise_pred, t, latents, generator=None): return (latents,)

class FakeMovq:
    config = SimpleNamespace(block_out_channels=[1, 1], latent_channels=4)

pipe = KandinskyV22Pipeline(FakeUNet(), FakeScheduler(), FakeMovq())
pipe.set_progress_bar_config(disable=True)
pipe(
    image_embeds=torch.randn(1, 32),
    negative_image_embeds=torch.randn(1, 32),
    num_images_per_prompt=2,
    guidance_scale=1.0,
    output_type="latent",
)

Relevant precedent:
The prior pipeline expands text conditioning before CFG concatenation:

Suggested fix:

image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0).to(
    dtype=self.unet.dtype, device=device
)

if self.do_classifier_free_guidance:
    negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0).to(
        dtype=self.unet.dtype, device=device
    )
    image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)

Apply the same pattern to hint in both ControlNet pipelines before duplicating it for CFG.

Issue 3: Inpainting batches preserve the first image/mask for every sample

Affected code:

Problem:
The inpaint loop uses image[:1] and mask_image[:1] when reinserting preserved regions. For batched inputs, every sample uses the first encoded image and first mask instead of its own.

Impact:
Batched inpainting silently produces wrong preserved regions. This can make all batch outputs inherit the first input image under preserved mask areas.

Reproduction:

from types import SimpleNamespace
import torch
from diffusers import KandinskyV22InpaintPipeline

class FakeUNet:
    dtype = torch.float32
    config = SimpleNamespace(in_channels=9)
    def __call__(self, sample, timestep, encoder_hidden_states=None, added_cond_kwargs=None, return_dict=False):
        return (torch.zeros(sample.shape[0], 8, sample.shape[2], sample.shape[3]),)

class FakeScheduler:
    init_noise_sigma = 1.0
    config = SimpleNamespace(variance_type="learned")
    def set_timesteps(self, *args, **kwargs): self.timesteps = torch.tensor([1])
    def step(self, noise_pred, t, latents, generator=None): return (latents,)

class FakeMovq:
    config = SimpleNamespace(block_out_channels=[1, 1], latent_channels=4)
    def encode(self, image):
        values = image.mean(dim=(1, 2, 3), keepdim=True)
        return {"latents": values.expand(image.shape[0], 4, 32, 32)}

pipe = KandinskyV22InpaintPipeline(FakeUNet(), FakeScheduler(), FakeMovq())
pipe.set_progress_bar_config(disable=True)

image = torch.stack([torch.full((3, 64, 64), -1.0), torch.full((3, 64, 64), 1.0)])
mask = torch.zeros(2, 1, 64, 64)  # preserve both inputs
out = pipe(
    image_embeds=torch.randn(2, 32),
    negative_image_embeds=torch.randn(2, 32),
    image=image,
    mask_image=mask,
    output_type="latent",
).images
print(torch.unique(out[0]).item(), torch.unique(out[1]).item())

Relevant precedent:
Stable Diffusion inpainting keeps per-sample mask and masked-image latents batched through the denoising loop rather than slicing to the first item.

Suggested fix:

init_latents = image.repeat_interleave(num_images_per_prompt, dim=0)
init_mask = mask_image.repeat_interleave(num_images_per_prompt, dim=0)

mask_image = init_mask
masked_image = init_latents * init_mask
if self.do_classifier_free_guidance:
    mask_image = mask_image.repeat(2, 1, 1, 1)
    masked_image = masked_image.repeat(2, 1, 1, 1)

...
init_latents_proper = init_latents
if i < len(timesteps) - 1:
    noise_timestep = timesteps[i + 1]
    init_latents_proper = self.scheduler.add_noise(init_latents, noise, noise_timestep[None])
latents = init_mask * init_latents_proper + (1 - init_mask) * latents

...
latents = init_mask * init_latents + (1 - init_mask) * latents

Issue 4: User-provided latents are not cast to the requested dtype

Affected code:

Problem:
The copied prepare_latents() blocks cast generated latents to dtype, but user-provided latents only call .to(device). In half precision, a float32 latent tensor stays float32 and can be passed into fp16 modules.

Impact:
Supplying reusable latents to fp16 pipelines can fail with dtype mismatch errors or force unintended float32 compute.

Reproduction:

from types import SimpleNamespace
import torch
from diffusers import KandinskyV22Pipeline

pipe = KandinskyV22Pipeline.__new__(KandinskyV22Pipeline)
scheduler = SimpleNamespace(init_noise_sigma=1.0)
latents = torch.randn(1, 4, 32, 32, dtype=torch.float32)

out = KandinskyV22Pipeline.prepare_latents(
    pipe, latents.shape, torch.float16, torch.device("cpu"), None, latents, scheduler
)
print(out.dtype)  # torch.float32, expected torch.float16

Relevant precedent:
Flux casts provided latents with both device and dtype:

Suggested fix:

latents = latents.to(device=device, dtype=dtype)

Because this is a copied block, either update the copied source and run make fix-copies, or remove the copy annotation if Kandinsky needs target-specific behavior.

Issue 5: Deprecated callback crashes when callback_steps is omitted

Affected code:

Problem:
The deprecated callback kwarg is still accepted, but callback_steps defaults to None after kwargs.pop("callback_steps", None). The loop then evaluates i % callback_steps, causing TypeError.

Impact:
Backward-compatible callback usage fails unless users also pass a deprecated callback_steps kwarg.

Reproduction:

from types import SimpleNamespace
import torch
from diffusers import KandinskyV22Pipeline

class FakeUNet:
    dtype = torch.float32
    config = SimpleNamespace(in_channels=4)
    def __call__(self, sample, timestep, encoder_hidden_states=None, added_cond_kwargs=None, return_dict=False):
        return (torch.zeros(sample.shape[0], 8, sample.shape[2], sample.shape[3]),)

class FakeScheduler:
    init_noise_sigma = 1.0
    config = SimpleNamespace(variance_type="learned")
    def set_timesteps(self, *args, **kwargs): self.timesteps = torch.tensor([1])
    def step(self, noise_pred, t, latents, generator=None): return (latents,)

class FakeMovq:
    config = SimpleNamespace(block_out_channels=[1, 1], latent_channels=4)

pipe = KandinskyV22Pipeline(FakeUNet(), FakeScheduler(), FakeMovq())
pipe.set_progress_bar_config(disable=True)
pipe(
    image_embeds=torch.randn(1, 32),
    negative_image_embeds=torch.randn(1, 32),
    output_type="latent",
    callback=lambda step, t, latents: None,
)

Relevant precedent:
The ControlNet variants still expose callback_steps: int = 1 in the signature:

Suggested fix:

callback_steps = kwargs.pop("callback_steps", 1)
if callback_steps is None:
    callback_steps = 1

Also validate that callback_steps is a positive integer before the denoising loop.

Issue 6: Slow coverage is missing for several exported pipelines

Affected code:

Problem:
The family has fast tests, and text/img2img/inpaint decoder pipelines have @slow integration tests. But KandinskyV22PriorPipeline, KandinskyV22PriorEmb2EmbPipeline, all three combined pipelines, and both ControlNet pipelines do not have @slow tests. The ControlNet integration tests are @nightly, which means they are not collected by slow CI.

Impact:
Published checkpoints and connected-pipeline loading paths can regress without slow-suite coverage. This also leaves the broken PriorEmb2Emb.interpolate() path untested.

Reproduction:

from pathlib import Path

for path in sorted(Path("tests/pipelines/kandinsky2_2").glob("test_*.py")):
    text = path.read_text()
    print(f"{path.name}: slow={'@slow' in text}, nightly={'@nightly' in text}")

Relevant precedent:
Existing slow coverage for decoder text/img2img/inpaint:

Suggested fix:
Add @slow integration tests for prior, prior emb2emb interpolation, combined text/img2img/inpaint loading, and ControlNet text/img2img. Keep @nightly if desired, but add @slow so slow CI covers the family.

Notes

Public exports, lazy loading, top-level imports, dummy objects, AutoPipeline registrations, docs, examples, and fast tests were checked. Top-level imports for all exported KandinskyV22* classes succeed locally.

I attempted .venv pytest collection, but this environment's torch build fails during test import with ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package, before Kandinsky tests are collected.