kandinsky3 model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Issue 1: Kandinsky3Img2ImgPipeline re-encodes latent image inputs

Affected code:

Problem:
VaeImageProcessor.preprocess() returns 4-channel latent tensors unchanged, but the pipeline then always calls self.movq.encode(image). A user-provided latent tensor is sent into a 3-channel image encoder and fails.

Impact:
Direct latent img2img inputs are unusable, and the latent branch in prepare_latents() is effectively unreachable from __call__.

Reproduction:

import torch
from diffusers import Kandinsky3Img2ImgPipeline, VQModel
from diffusers.schedulers import DDPMScheduler

movq = VQModel(block_out_channels=[32], down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"], in_channels=3, out_channels=3, latent_channels=4, layers_per_block=1, norm_num_groups=8, num_vq_embeddings=12, vq_embed_dim=4)
pipe = Kandinsky3Img2ImgPipeline(None, None, None, DDPMScheduler(num_train_timesteps=4), movq)
pipe(prompt_embeds=torch.ones(1, 2, 4), attention_mask=torch.ones(1, 2, dtype=torch.long), image=torch.randn(1, 4, 8, 8), guidance_scale=1.0, num_inference_steps=1, output_type="latent")

Relevant precedent:
VaeImageProcessor.preprocess() intentionally returns latent-channel tensors unchanged:

Suggested fix:

if image.shape[1] == self.movq.config.latent_channels:
    latents = image
else:
    latents = self.movq.encode(image)["latents"]
latents = latents.repeat_interleave(num_images_per_prompt, dim=0)

Issue 2: strength=0.0 returns an empty batch

Affected code:

Problem:
The docs say strength is between 0 and 1, but strength=0.0 produces an empty timestep tensor. That empty tensor is used as latent_timestep, and the pipeline returns an empty latent batch instead of preserving the input image or raising.

Impact:
A boundary value accepted by the public API silently returns torch.Size([0, ...]).

Reproduction:

import torch
from diffusers import Kandinsky3Img2ImgPipeline, VQModel
from diffusers.schedulers import DDPMScheduler

movq = VQModel(block_out_channels=[32], down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"], in_channels=3, out_channels=3, latent_channels=4, layers_per_block=1, norm_num_groups=8, num_vq_embeddings=12, vq_embed_dim=4)
pipe = Kandinsky3Img2ImgPipeline(None, None, None, DDPMScheduler(num_train_timesteps=4), movq)
out = pipe(prompt_embeds=torch.ones(1, 2, 4), attention_mask=torch.ones(1, 2, dtype=torch.long), image=torch.rand(1, 3, 8, 8), strength=0.0, guidance_scale=1.0, num_inference_steps=2, output_type="latent")
print(out.images.shape)  # torch.Size([0, 4, 8, 8])

Relevant precedent:
Stable Diffusion img2img validates the public strength range:

Suggested fix:
Handle strength == 0 explicitly by returning the initial latents/image without denoising, or reject it:

if strength <= 0 or strength > 1:
    raise ValueError(f"The value of strength should be in (0.0, 1.0], but is {strength}")

Issue 3: Kandinsky3UNet crashes when encoder_attention_mask is omitted

Affected code:

Problem:
encoder_attention_mask defaults to None, but Kandinsky3AttentionPooling.forward() unconditionally calls context_mask.to(...).

Impact:
Direct model users cannot rely on the optional mask default, and model-level tests do not catch it because the pipelines always pass masks.

Reproduction:

import torch
from diffusers import Kandinsky3UNet

m = Kandinsky3UNet(in_channels=4, time_embedding_dim=4, groups=2, attention_head_dim=4, layers_per_block=1, block_out_channels=(32, 64), cross_attention_dim=4, encoder_hid_dim=32)
m(torch.randn(1, 4, 8, 8), torch.tensor(1), encoder_hidden_states=torch.randn(1, 2, 32), return_dict=False)

Relevant precedent:
Generic attention processors accept attention_mask=None.

Suggested fix:

def forward(self, x, context, context_mask=None):
    if context_mask is not None:
        context_mask = context_mask.to(dtype=context.dtype)
    context = self.attention(context.mean(dim=1, keepdim=True), context, context_mask)
    return x + context.squeeze(1)

Issue 4: Tuple-valued UNet config annotations are not implemented

Affected code:

Problem:
layers_per_block and cross_attention_dim are typed as accepting tuples, but tuple values fail during construction. layers_per_block is copied as a tuple into every level, and cross_attention_dim is passed to nn.Linear as a tuple.

Impact:
The serialized/public config surface advertises values that cannot be loaded.

Reproduction:

from diffusers import Kandinsky3UNet

base = dict(in_channels=4, time_embedding_dim=4, groups=2, attention_head_dim=4, block_out_channels=(32, 64), encoder_hid_dim=32)
Kandinsky3UNet(**base, layers_per_block=(1, 1), cross_attention_dim=4)
Kandinsky3UNet(**base, layers_per_block=1, cross_attention_dim=(4, 4))

Relevant precedent:
UNet2DConditionModel expands scalar-or-tuple config fields before constructing blocks:

Suggested fix:
Support tuple layers_per_block, and either implement per-block context projections or reject tuple cross_attention_dim explicitly:

if isinstance(layers_per_block, int):
    num_blocks = [layers_per_block] * len(block_out_channels)
else:
    if len(layers_per_block) != len(block_out_channels):
        raise ValueError("`layers_per_block` must match `block_out_channels`.")
    num_blocks = list(layers_per_block)

if not isinstance(cross_attention_dim, int):
    raise ValueError("`Kandinsky3UNet` currently supports only an integer `cross_attention_dim`.")

Issue 5: Deprecated callback crashes when callback_steps is omitted

Affected code:

Problem:
callback_steps defaults to None, but the loop evaluates i % callback_steps whenever deprecated callback is passed.

Impact:
The deprecated callback API is still present until 1.0.0 but fails unless users also know to pass deprecated callback_steps.

Reproduction:

import torch
from diffusers import Kandinsky3Pipeline
from diffusers.schedulers import DDPMScheduler

class ZeroUNet(torch.nn.Module):
    @property
    def dtype(self): return torch.float32
    @property
    def device(self): return torch.device("cpu")
    def forward(self, sample, *args, **kwargs): return (torch.zeros_like(sample),)

pipe = Kandinsky3Pipeline(None, None, ZeroUNet(), DDPMScheduler(num_train_timesteps=4), None)
pipe(prompt_embeds=torch.ones(1, 2, 4), attention_mask=torch.ones(1, 2, dtype=torch.long), guidance_scale=1.0, height=8, width=8, num_inference_steps=1, output_type="latent", callback=lambda step, timestep, latents: None)

Relevant precedent:
Deprecated arguments should continue to work until the removal version.

Suggested fix:

callback = kwargs.pop("callback", None)
callback_steps = kwargs.pop("callback_steps", 1 if callback is not None else None)

Issue 6: encode_prompt() is decorated with @torch.no_grad()

Affected code:

Problem:
The review rules call out helper-level @torch.no_grad() as incorrect because __call__ already owns inference no-grad behavior. Keeping it on encode_prompt() blocks advanced callers from using gradients for prompt embedding workflows.

Impact:
Public helper behavior is less flexible than related modern pipelines.

Reproduction:

import torch
from types import SimpleNamespace
from diffusers import Kandinsky3Pipeline

class Tokenizer:
    def __call__(self, prompt, max_length=None, **kwargs):
        return SimpleNamespace(input_ids=torch.arange(max_length).unsqueeze(0), attention_mask=torch.ones(1, max_length, dtype=torch.long))

class TextEncoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = torch.nn.Embedding(128, 4)
    @property
    def dtype(self): return self.emb.weight.dtype
    def forward(self, input_ids, attention_mask=None): return (self.emb(input_ids),)

pipe = Kandinsky3Pipeline(Tokenizer(), TextEncoder(), None, None, None)
print(pipe.encode_prompt("x", do_classifier_free_guidance=False, device="cpu")[0].requires_grad)  # False

Relevant precedent:
FluxPipeline.encode_prompt() is not decorated, while __call__ is:

Suggested fix:
Remove @torch.no_grad() from both encode_prompt() methods.

Issue 7: Kandinsky3 conversion script constructs the UNet with a positional config dict

Affected code:

Problem:
Kandinsky3UNet(config) passes the dict as in_channels, so construction fails inside convolution/group norm setup.

Impact:
The checked-in converter cannot run from its CLI path.

Reproduction:

from diffusers import Kandinsky3UNet
Kandinsky3UNet({})

Relevant precedent:
Top-level model constructors should receive config values as keyword arguments.

Suggested fix:

config = {}
unet = Kandinsky3UNet(**config)
unet.load_state_dict(converted_state_dict, strict=True)

Coverage and duplicate-search status

Public imports, lazy loading, auto-pipeline mappings, dummy objects, pipeline runtime paths, UNet config/runtime behavior, docs, fast tests, and slow tests were reviewed.

Slow tests are present for text2image and img2img in tests/pipelines/kandinsky3/. There are no standalone tests/models tests for Kandinsky3UNet; coverage is pipeline-only.

Attempted fast pytest targets in .venv, but collection failed before the tests ran because this Torch install lacks torch._C._distributed_c10d.

Duplicate search performed with gh search issues --include-prs against huggingface/diffusers for kandinsky3, Kandinsky3UNet, convert_kandinsky3_unet, Kandinsky3Img2ImgPipeline strength, callback_steps, latent image, encode_prompt no_grad, layers_per_block, encoder_attention_mask, and context_mask. I found related historical items, including #5963, #11080, #12474, and #12544, but no direct duplicate for the issues above.