kandinsky model/pipeline review

kandinsky model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Duplicate search performed with gh search issues and gh search prs on huggingface/diffusers for kandinsky, affected class names, image_embeds, num_images_per_prompt, combined PIL inputs, transformer_kandinsky, bfloat16, device_map, _no_split_modules, and slow-test coverage. No likely duplicates were found.

Issue 1: Decoder image embeddings are only expanded in the CFG path

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py#L330-L336
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py#L521-L527

Problem:
image_embeds is repeated to num_images_per_prompt and moved to the execution dtype/device only when guidance_scale > 1. With guidance_scale <= 1, prompt_embeds and latents are expanded but image_embeds stays at the original batch size. KandinskyPipeline and KandinskyInpaintPipeline then fail for num_images_per_prompt > 1 unless the caller manually pre-expands image_embeds.

Impact:
Valid no-CFG calls cannot generate multiple images per prompt from one prior embedding. The no-CFG path also skips dtype/device normalization for image_embeds.

Reproduction:

import torch
from types import SimpleNamespace
from diffusers import DDIMScheduler, KandinskyPipeline

class Module(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.dummy = torch.nn.Parameter(torch.zeros(()))
    @property
    def dtype(self): return self.dummy.dtype
    @property
    def device(self): return self.dummy.device

class Tokenizer:
    model_max_length = 77
    def __call__(self, prompt, **kwargs):
        b = len(prompt) if isinstance(prompt, list) else 1
        return SimpleNamespace(input_ids=torch.ones(b, 4, dtype=torch.long), attention_mask=torch.ones(b, 4, dtype=torch.long))
    def batch_decode(self, ids): return [""]

class TextEncoder(Module):
    def forward(self, input_ids, attention_mask):
        b, s = input_ids.shape
        return torch.zeros(b, 32), torch.zeros(b, s, 32)

class UNet(Module):
    config = SimpleNamespace(in_channels=4)
    def forward(self, sample, timestep, encoder_hidden_states, added_cond_kwargs, 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], 8, sample.shape[2], sample.shape[3]),)

class Movq(Module):
    config = SimpleNamespace(block_out_channels=[1])
    def decode(self, latents, force_not_quantize=True): return {"sample": latents[:, :3]}

pipe = KandinskyPipeline(TextEncoder(), Tokenizer(), UNet(), DDIMScheduler(num_train_timesteps=2), Movq()).to("cpu")
pipe.set_progress_bar_config(disable=True)
pipe(prompt="p", image_embeds=torch.zeros(1, 32), negative_image_embeds=torch.zeros(1, 32), guidance_scale=1.0, num_images_per_prompt=2, num_inference_steps=1, output_type="pt")

Relevant precedent:
prompt_embeds is expanded before the CFG branch in the same methods:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py#L171-L173

Suggested fix:

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

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

Issue 2: Combined img2img/inpaint wrappers check prompt instead of image for PIL wrapping

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py#L537-L548
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py#L777-L796

Problem:
KandinskyImg2ImgCombinedPipeline and KandinskyInpaintCombinedPipeline use isinstance(prompt, PIL.Image.Image) when they mean to check image. A single PIL image is not wrapped into a list, so the subsequent expansion to match num_images_per_prompt is skipped.

Impact:
Combined img2img/inpaint calls with a single PIL image and num_images_per_prompt > 1 fail in the decoder with latent batch shape mismatches. Passing [image] works, which makes this a wrapper bug rather than a decoder limitation.

Reproduction:

import numpy as np
import torch
from PIL import Image
from types import SimpleNamespace
from diffusers import KandinskyImg2ImgCombinedPipeline

class Prior:
    def __call__(self, **kwargs):
        return (torch.zeros(2, 32), torch.zeros(2, 32))

class Decoder:
    def __call__(self, prompt, image, **kwargs):
        assert isinstance(image, list) and len(image) == 2, (type(image), image if isinstance(image, list) else None)
        return SimpleNamespace(images="ok")

pipe = object.__new__(KandinskyImg2ImgCombinedPipeline)
pipe.prior_pipe = Prior()
pipe.decoder_pipe = Decoder()
pipe.maybe_free_model_hooks = lambda: None

image = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8))
pipe(prompt="horse", image=image, num_images_per_prompt=2)

Relevant precedent:
The branch immediately below already handles list/tuple inputs by length; the missing step is wrapping the single PIL image first.

Suggested fix:

prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt
image = [image] if isinstance(image, PIL.Image.Image) else image

For inpaint, apply the same correction to image and keep the existing mask_image PIL wrapping.

Issue 3: Kandinsky5 transformer fails after direct half/bfloat16 casting

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_kandinsky.py#L168-L171
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_kandinsky.py#L309-L312

Problem:
Kandinsky5TimeEmbeddings.forward() always builds a float32 embedding and feeds it into nn.Linear. If users call model.to(torch.float16) or model.to(torch.bfloat16), the linear weights are converted but the input remains float32, causing a dtype mismatch. The rotary helper also hard-casts through torch.bfloat16, which silently quantizes float32/float16 attention states.

Impact:
Direct dtype conversion, a common Diffusers usage pattern, breaks Kandinsky5 transformer inference. The hard-coded rotary bfloat16 cast also prevents clean dtype behavior and can degrade fp32 parity.

Reproduction:

import torch
from diffusers import Kandinsky5Transformer3DModel

model = Kandinsky5Transformer3DModel(
    in_visual_dim=4, in_text_dim=8, in_text_dim2=4, time_dim=8, out_visual_dim=4,
    patch_size=(1, 1, 1), model_dim=8, ff_dim=16, num_text_blocks=1,
    num_visual_blocks=1, axes_dims=(2, 2, 4), visual_cond=False,
).eval().to(dtype=torch.bfloat16)

model(
    hidden_states=torch.randn(1, 1, 2, 2, 4, dtype=torch.bfloat16),
    encoder_hidden_states=torch.randn(1, 3, 8, dtype=torch.bfloat16),
    timestep=torch.tensor([1], dtype=torch.bfloat16),
    pooled_projections=torch.randn(1, 4, dtype=torch.bfloat16),
    visual_rope_pos=[torch.arange(1), torch.arange(2), torch.arange(2)],
    text_rope_pos=torch.arange(3),
    return_dict=False,
)

Relevant precedent:
Wan casts timestep projections to the time embedder parameter dtype before the linear and casts the result back to the hidden-state dtype:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_wan.py#L341-L344

Suggested fix:

time_embedder_dtype = self.in_layer.weight.dtype
if time_embed.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
    time_embed = time_embed.to(time_embedder_dtype)
time_embed = self.out_layer(self.activation(self.in_layer(time_embed)))

And for rotary:

orig_dtype = x.dtype
x_ = x.reshape(*x.shape[:-1], -1, 1, 2).float()
x_out = (rope * x_).sum(dim=-1)
return x_out.reshape(*x.shape).to(orig_dtype)

Issue 4: Kandinsky5 transformer declares repeated blocks but no _no_split_modules

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_kandinsky.py#L522-L527

Problem:
Kandinsky5Transformer3DModel sets _repeated_blocks but does not set _no_split_modules. Related transformer families mark block classes as no-split for device-map/offload placement.

Impact:
Automatic device maps can split transformer blocks across devices. That is inconsistent with related transformer models and risks poor placement, extra transfers, or correctness issues in block-local attention/modulation paths.

Reproduction:

from diffusers import Kandinsky5Transformer3DModel

print(getattr(Kandinsky5Transformer3DModel, "_no_split_modules", None))
print(Kandinsky5Transformer3DModel._repeated_blocks)
assert Kandinsky5Transformer3DModel._no_split_modules is not None

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L566-L568
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_qwenimage.py#L779-L781
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_wan.py#L548-L551

Suggested fix:

_no_split_modules = ["Kandinsky5TransformerEncoderBlock", "Kandinsky5TransformerDecoderBlock"]
_repeated_blocks = ["Kandinsky5TransformerEncoderBlock", "Kandinsky5TransformerDecoderBlock"]

Issue 5: Slow and enforced fast coverage is incomplete

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky.py#L220-L225
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky.py#L285-L312
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky_inpaint.py#L236-L241
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky_inpaint.py#L307-L323
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky_combined.py#L76-L81
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky_prior.py#L171-L240

Problem:
The core fast image-slice tests are xfailed for Transformers >=4.56.2, so they are not enforced on the current test stack. Slow coverage exists for text2img and img2img, but not for KandinskyPriorPipeline or the combined pipelines. Inpaint has a nightly integration test but no @slow test. There is no dedicated tests/models coverage for Kandinsky5Transformer3DModel.

Impact:
The regressions above are not caught: no-CFG multi-image decoder behavior, combined PIL expansion, Kandinsky5 dtype conversion, and device-map metadata all slip through the current suite.

Reproduction:

from pathlib import Path

for path in sorted(Path("tests/pipelines/kandinsky").glob("test_*.py")):
    text = path.read_text()
    print(path.as_posix(), "slow=", "@slow" in text, "nightly=", "@nightly" in text, "xfail=", "xfail" in text)

has_model_test = any(
    "Kandinsky5Transformer3DModel" in path.read_text(errors="ignore")
    for path in Path("tests/models").rglob("test_*.py")
)
print("Kandinsky5 model test:", has_model_test)

Relevant precedent:
The existing text2img and img2img files show the expected slow integration-test shape:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky.py#L285-L312
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky_img2img.py#L309-L372

Suggested fix:
Add or re-enable coverage for:

# Fast
# - no-CFG num_images_per_prompt > 1 for KandinskyPipeline and KandinskyInpaintPipeline
# - single PIL image + num_images_per_prompt > 1 for combined img2img/inpaint
# - Kandinsky5Transformer3DModel .to(torch.bfloat16) forward

# Slow
# - KandinskyPriorPipeline
# - KandinskyCombinedPipeline
# - KandinskyImg2ImgCombinedPipeline
# - KandinskyInpaintCombinedPipeline
# - @slow inpaint coverage, or mark the existing nightly integration as slow too

Update the Transformers >=4.56.2 expected slices instead of leaving the main fast output tests xfailed.

Start with the affected pipeline files under src/diffusers/pipelines/kandinsky and transformer_kandinsky.py, then run the reproductions in the issue. Review the referenced tests in tests/pipelines/kandinsky and add coverage for no-CFG image expansion, combined PIL inputs, Kandinsky5 dtype/device-map behavior, and missing slow tests. Done means the reproductions pass and the relevant regressions are enforced by tests.