kandinsky3 model/pipeline review

[hlky]

kandinsky3 model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Issue 1: Kandinsky3Img2ImgPipeline re-encodes latent image inputs

Affected code:

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py

Lines 559 to 565 in 0f1abc4

	image = torch.cat([self.image_processor.preprocess(i) for i in image], dim=0)
	image = image.to(dtype=prompt_embeds.dtype, device=device)
	# 4. Prepare timesteps
	self.scheduler.set_timesteps(num_inference_steps, device=device)
	timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
	# 5. Prepare latents
	latents = self.movq.encode(image)["latents"]

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:

diffusers/src/diffusers/image_processor.py

Lines 712 to 715 in 0f1abc4

	channel = image.shape[1]
	# don't need any preprocess if the image is latents
	if channel == self.config.vae_latent_channels:
	return image

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:

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py

Lines 143 to 150 in 0f1abc4

	attention_mask (`torch.Tensor`, *optional*):
	Pre-generated attention mask. Must provide if passing `prompt_embeds` directly.
	negative_attention_mask (`torch.Tensor`, *optional*):
	Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly.
	"""
	if prompt is not None and negative_prompt is not None:
	if type(prompt) is not type(negative_prompt):
	raise TypeError(

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py

Lines 562 to 568 in 0f1abc4

	self.scheduler.set_timesteps(num_inference_steps, device=device)
	timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
	# 5. Prepare latents
	latents = self.movq.encode(image)["latents"]
	latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
	latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
	latents = self.prepare_latents(

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:

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

Lines 656 to 669 in 0f1abc4

	def check_inputs(
	self,
	prompt,
	strength,
	callback_steps,
	negative_prompt=None,
	prompt_embeds=None,
	negative_prompt_embeds=None,
	ip_adapter_image=None,
	ip_adapter_image_embeds=None,
	callback_on_step_end_tensor_inputs=None,
	):
	if strength < 0 or strength > 1:
	raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")

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:

diffusers/src/diffusers/models/unets/unet_kandinsky3.py

Lines 149 to 152 in 0f1abc4

	def forward(self, sample, timestep, encoder_hidden_states=None, encoder_attention_mask=None, return_dict=True):
	if encoder_attention_mask is not None:
	encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
	encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

diffusers/src/diffusers/models/unets/unet_kandinsky3.py

Lines 431 to 433 in 0f1abc4

	def forward(self, x, context, context_mask=None):
	context_mask = context_mask.to(dtype=context.dtype)
	context = self.attention(context.mean(dim=1, keepdim=True), context, context_mask)

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:

diffusers/src/diffusers/models/unets/unet_kandinsky3.py

Lines 56 to 58 in 0f1abc4

	layers_per_block: int | tuple[int] = 3,
	block_out_channels: tuple[int, ...] = (384, 768, 1536, 3072),
	cross_attention_dim: int | tuple[int] = 4096,

diffusers/src/diffusers/models/unets/unet_kandinsky3.py

Lines 88 to 90 in 0f1abc4

	text_dims = [cross_attention_dim if is_exist else None for is_exist in add_cross_attention]
	num_blocks = len(block_out_channels) * [layers_per_block]
	layer_params = [num_blocks, text_dims, add_self_attention]

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:

diffusers/src/diffusers/models/unets/unet_2d_condition.py

Lines 337 to 344 in 0f1abc4

	if isinstance(num_attention_heads, int):
	num_attention_heads = (num_attention_heads,) * len(down_block_types)
	if isinstance(attention_head_dim, int):
	attention_head_dim = (attention_head_dim,) * len(down_block_types)
	if isinstance(cross_attention_dim, int):
	cross_attention_dim = (cross_attention_dim,) * len(down_block_types)

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:

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py

Lines 430 to 431 in 0f1abc4

	callback = kwargs.pop("callback", None)
	callback_steps = kwargs.pop("callback_steps", None)

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py

Lines 555 to 559 in 0f1abc4

	if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
	progress_bar.update()
	if callback is not None and i % callback_steps == 0:
	step_idx = i // getattr(self.scheduler, "order", 1)
	callback(step_idx, t, latents)

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py

Lines 488 to 489 in 0f1abc4

	callback = kwargs.pop("callback", None)
	callback_steps = kwargs.pop("callback_steps", None)

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py

Lines 613 to 617 in 0f1abc4

	if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
	progress_bar.update()
	if callback is not None and i % callback_steps == 0:
	step_idx = i // getattr(self.scheduler, "order", 1)
	callback(step_idx, t, latents)

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:

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py

Lines 91 to 92 in 0f1abc4

	@torch.no_grad()
	def encode_prompt(

diffusers/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py

Lines 106 to 107 in 0f1abc4

	@torch.no_grad()
	def encode_prompt(

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:

diffusers/src/diffusers/pipelines/flux/pipeline_flux.py

Lines 311 to 316 in 0f1abc4

	def encode_prompt(
	self,
	prompt: str | list[str],
	prompt_2: str | list[str] | None = None,
	device: torch.device | None = None,
	num_images_per_prompt: int = 1,

diffusers/src/diffusers/pipelines/flux/pipeline_flux.py

Lines 652 to 654 in 0f1abc4

	@torch.no_grad()
	@replace_example_docstring(EXAMPLE_DOC_STRING)
	def __call__(

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:

diffusers/src/diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py

Lines 79 to 86 in 0f1abc4

	# Initialize your Kandinsky3UNet model
	config = {}
	# Convert the state dict
	converted_state_dict = convert_state_dict(unet_state_dict)
	unet = Kandinsky3UNet(config)
	unet.load_state_dict(converted_state_dict)

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.