huggingface / huggingface/diffusers
kandinsky2_2 model/pipeline review
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
# `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: https://github.com/huggingface/diffusers/issues/4183 covers the ControlNet `guidance_scale <= 1` subset of Issue 2; https://github.com/huggingface/diffusers/issues/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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py#L220-L230
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py#L399-L413
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py#L244-L246
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py#L197-L241
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py#L227-L239
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py#L286-L298
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py#L420-L432
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py#L236-L253
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py#L278-L295
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 https://github.com/huggingface/diffusers/issues/4183.
Reproduction:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py#L301-L303
Suggested fix:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py#L510-L519
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py#L540-L541
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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py#L106-L115
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py#L243-L252
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py#L279-L287
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py#L149-L157
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/pipeline_flux.py#L615-L618
Suggested fix:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py#L200-L214
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py#L303-L305
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py#L259-L273
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py#L393-L407
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py#L171-L176
Suggested fix:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py#L199-L245
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py#L201-L238
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py#L224-L246
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet_img2img.py#L230-L252
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py#L37-L58
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky.py#L225-L240
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky_img2img.py#L242-L257
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py#L303-L318
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.
Contributor guide
Research direction
Start with the affected entry points in the kandinsky2_2 pipeline files, especially pipeline_kandinsky2_2_prior_emb2emb.py, pipeline_kandinsky2_2.py, pipeline_kandinsky2_2_img2img.py, pipeline_kandinsky2_2_inpainting.py, and the ControlNet variants. Run the reproductions in the issue and compare the corresponding prior pipeline and Flux implementations. Done means the five reported behaviors are covered without regressions across batching, no-CFG execution, dtype handling, interpolation, and callbacks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100