huggingface / huggingface/diffusers

kandinsky2_2 model/pipeline review

Aperta
#13,596 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

# `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.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Inizia dai punti di ingresso interessati nei file della pipeline di kandinsky2_2, in particolare pipeline_kandinsky2_2_prior_emb2emb.py, pipeline_kandinsky2_2.py, pipeline_kandinsky2_2_img2img.py, pipeline_kandinsky2_2_inpainting.py e le varianti di ControlNet. Esegui le repro dell’issue e confronta la prior pipeline corrispondente con le implementazioni di Flux. Il lavoro è completato quando i cinque comportamenti segnalati sono coperti senza regressioni per batching, esecuzione no-CFG, gestione del dtype, interpolazione e callbacks.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
5/5
Tempo stimato
Più di una settimana
Stato di attività
Tranquilla
Chiarezza
Abbastanza chiara
Idoneità per principianti
42/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.