Skip to content

Fix batching and precomputed embeddings in Kandinsky 5 pipelines #14168

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
pzarzycki wants to merge 1 commit into
base: main
Choose a base branch
from
+308 −80
Draft

Fix batching and precomputed embeddings in Kandinsky 5 pipelines #14168

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 36 additions & 18 deletions src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -456,8 +456,6 @@ def encode_prompt(
if not isinstance(prompt, list):
prompt = [prompt]

batch_size = len(prompt)

prompt = [prompt_clean(p) for p in prompt]

# Encode with Qwen2.5-VL
Expand All @@ -479,20 +477,10 @@ def encode_prompt(

# Repeat embeddings for num_videos_per_prompt
# Qwen embeddings: repeat sequence for each video, then reshape
prompt_embeds_qwen = prompt_embeds_qwen.repeat(
1, num_videos_per_prompt, 1
) # [batch_size, seq_len * num_videos_per_prompt, embed_dim]
# Reshape to [batch_size * num_videos_per_prompt, seq_len, embed_dim]
prompt_embeds_qwen = prompt_embeds_qwen.view(
batch_size * num_videos_per_prompt, -1, prompt_embeds_qwen.shape[-1]
)
prompt_embeds_qwen = prompt_embeds_qwen.repeat_interleave(num_videos_per_prompt, dim=0)

# CLIP embeddings: repeat for each video
prompt_embeds_clip = prompt_embeds_clip.repeat(
1, num_videos_per_prompt, 1
) # [batch_size, num_videos_per_prompt, clip_embed_dim]
# Reshape to [batch_size * num_videos_per_prompt, clip_embed_dim]
prompt_embeds_clip = prompt_embeds_clip.view(batch_size * num_videos_per_prompt, -1)
prompt_embeds_clip = prompt_embeds_clip.repeat_interleave(num_videos_per_prompt, dim=0)

# Repeat cumulative sequence lengths for num_videos_per_prompt
# Original cu_seqlens: [0, len1, len1+len2, ...]
Expand Down Expand Up @@ -565,6 +553,13 @@ def check_inputs(
"If any of `prompt_embeds_qwen`, `prompt_embeds_clip`, or `prompt_cu_seqlens` is provided, "
"all three must be provided."
)
if (
prompt_embeds_qwen.shape[0] != prompt_embeds_clip.shape[0]
or prompt_cu_seqlens.shape[0] != prompt_embeds_qwen.shape[0] + 1
):
raise ValueError(
"Precomputed positive prompt embeddings and sequence lengths must have matching batch sizes."
)

# Check for consistency within negative prompt embeddings and sequence lengths
if (
Expand All @@ -581,6 +576,13 @@ def check_inputs(
"If any of `negative_prompt_embeds_qwen`, `negative_prompt_embeds_clip`, or `negative_prompt_cu_seqlens` is provided, "
"all three must be provided."
)
if (
negative_prompt_embeds_qwen.shape[0] != negative_prompt_embeds_clip.shape[0]
or negative_prompt_cu_seqlens.shape[0] != negative_prompt_embeds_qwen.shape[0] + 1
):
raise ValueError(
"Precomputed negative prompt embeddings and sequence lengths must have matching batch sizes."
)

# Check if prompt or embeddings are provided (either prompt or all required embedding components for positive)
if prompt is None and prompt_embeds_qwen is None:
Expand Down Expand Up @@ -806,31 +808,47 @@ def __call__(
if prompt_embeds_qwen is None:
prompt_embeds_qwen, prompt_embeds_clip, prompt_cu_seqlens = self.encode_prompt(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
else:
prompt_embeds_qwen = prompt_embeds_qwen.repeat_interleave(num_videos_per_prompt, dim=0)
prompt_embeds_clip = prompt_embeds_clip.repeat_interleave(num_videos_per_prompt, dim=0)
prompt_lengths = prompt_cu_seqlens.diff().repeat_interleave(num_videos_per_prompt)
prompt_cu_seqlens = F.pad(prompt_lengths.cumsum(0), (1, 0), value=0)

if self.guidance_scale > 1.0:
if negative_prompt is None:
if negative_prompt is None and negative_prompt_embeds_qwen is None:
negative_prompt = "Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, deformed, walking backwards"

if isinstance(negative_prompt, str):
negative_prompt = [negative_prompt] * len(prompt) if prompt is not None else [negative_prompt]
elif len(negative_prompt) != len(prompt):
negative_prompt = [negative_prompt] * batch_size
elif negative_prompt is not None and len(negative_prompt) != batch_size:
raise ValueError(
f"`negative_prompt` must have same length as `prompt`. Got {len(negative_prompt)} vs {len(prompt)}."
f"`negative_prompt` must have the same batch size as `prompt`. Got {len(negative_prompt)} vs {batch_size}."
)

if negative_prompt_embeds_qwen is None:
negative_prompt_embeds_qwen, negative_prompt_embeds_clip, negative_prompt_cu_seqlens = (
self.encode_prompt(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
)
else:
negative_prompt_embeds_qwen = negative_prompt_embeds_qwen.repeat_interleave(
num_videos_per_prompt, dim=0
)
negative_prompt_embeds_clip = negative_prompt_embeds_clip.repeat_interleave(
num_videos_per_prompt, dim=0
)
negative_prompt_lengths = negative_prompt_cu_seqlens.diff().repeat_interleave(num_videos_per_prompt)
negative_prompt_cu_seqlens = F.pad(negative_prompt_lengths.cumsum(0), (1, 0), value=0)

# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
Expand Down
71 changes: 48 additions & 23 deletions src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,8 +207,13 @@ def _encode_prompt_qwen(
"""
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
if not isinstance(image, list):
image = [image]
image = self.image_processor.postprocess(self.image_processor.preprocess(image), output_type="pil")
if len(image) == 1:
image = image * len(prompt)
elif len(image) != len(prompt):
raise ValueError(
f"The image batch size must be 1 or match the prompt batch size, but got {len(image)} and {len(prompt)}."
)
image = [i.resize((i.size[0] // 2, i.size[1] // 2)) for i in image]
full_texts = [self.prompt_template.format(p) for p in prompt]
max_allowed_len = self.prompt_template_encode_start_idx + max_sequence_length
Expand Down Expand Up @@ -332,8 +337,6 @@ def encode_prompt(
if not isinstance(prompt, list):
prompt = [prompt]

batch_size = len(prompt)

prompt = [prompt_clean(p) for p in prompt]

# Encode with Qwen2.5-VL
Expand All @@ -356,20 +359,10 @@ def encode_prompt(

# Repeat embeddings for num_images_per_prompt
# Qwen embeddings: repeat sequence for each image, then reshape
prompt_embeds_qwen = prompt_embeds_qwen.repeat(
1, num_images_per_prompt, 1
) # [batch_size, seq_len * num_images_per_prompt, embed_dim]
# Reshape to [batch_size * num_images_per_prompt, seq_len, embed_dim]
prompt_embeds_qwen = prompt_embeds_qwen.view(
batch_size * num_images_per_prompt, -1, prompt_embeds_qwen.shape[-1]
)
prompt_embeds_qwen = prompt_embeds_qwen.repeat_interleave(num_images_per_prompt, dim=0)

# CLIP embeddings: repeat for each image
prompt_embeds_clip = prompt_embeds_clip.repeat(
1, num_images_per_prompt, 1
) # [batch_size, num_images_per_prompt, clip_embed_dim]
# Reshape to [batch_size * num_images_per_prompt, clip_embed_dim]
prompt_embeds_clip = prompt_embeds_clip.view(batch_size * num_images_per_prompt, -1)
prompt_embeds_clip = prompt_embeds_clip.repeat_interleave(num_images_per_prompt, dim=0)

# Repeat cumulative sequence lengths for num_images_per_prompt
# Original differences (lengths) for each prompt in the batch
Expand Down Expand Up @@ -448,6 +441,13 @@ def check_inputs(
"If any of `prompt_embeds_qwen`, `prompt_embeds_clip`, or `prompt_cu_seqlens` is provided, "
"all three must be provided."
)
if (
prompt_embeds_qwen.shape[0] != prompt_embeds_clip.shape[0]
or prompt_cu_seqlens.shape[0] != prompt_embeds_qwen.shape[0] + 1
):
raise ValueError(
"Precomputed positive prompt embeddings and sequence lengths must have matching batch sizes."
)

# Check for consistency within negative prompt embeddings and sequence lengths
if (
Expand All @@ -464,6 +464,13 @@ def check_inputs(
"If any of `negative_prompt_embeds_qwen`, `negative_prompt_embeds_clip`, or `negative_prompt_cu_seqlens` is provided, "
"all three must be provided."
)
if (
negative_prompt_embeds_qwen.shape[0] != negative_prompt_embeds_clip.shape[0]
or negative_prompt_cu_seqlens.shape[0] != negative_prompt_embeds_qwen.shape[0] + 1
):
raise ValueError(
"Precomputed negative prompt embeddings and sequence lengths must have matching batch sizes."
)

# Check if prompt or embeddings are provided (either prompt or all required embedding components for positive)
if prompt is None and prompt_embeds_qwen is None:
Expand Down Expand Up @@ -533,6 +540,11 @@ def prepare_latents(
# Encode the input image to use as first frame
# Preprocess image
image_tensor = self.image_processor.preprocess(image, height=height, width=width).to(device, dtype=dtype)
if batch_size % image_tensor.shape[0] != 0:
raise ValueError(
f"The effective batch size {batch_size} must be divisible by the image batch size {image_tensor.shape[0]}."
)
image_tensor = image_tensor.repeat_interleave(batch_size // image_tensor.shape[0], dim=0)
# Encode image to latents using VAE
with torch.no_grad():
image_latents = self.vae.encode(image_tensor).latent_dist.sample(generator=generator)
Expand Down Expand Up @@ -649,7 +661,7 @@ def __call__(
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
# 1. Check inputs. Raise error if not correct
if height is None and width is None:
width, height = image[0].size if isinstance(image, list) else image.size
height, width = self.image_processor.get_default_height_width(image)
self.check_inputs(
prompt=prompt,
negative_prompt=negative_prompt,
Expand Down Expand Up @@ -695,18 +707,22 @@ def __call__(
device=device,
dtype=dtype,
)
else:
prompt_embeds_qwen = prompt_embeds_qwen.repeat_interleave(num_images_per_prompt, dim=0)
prompt_embeds_clip = prompt_embeds_clip.repeat_interleave(num_images_per_prompt, dim=0)
prompt_lengths = prompt_cu_seqlens.diff().repeat_interleave(num_images_per_prompt)
prompt_cu_seqlens = F.pad(prompt_lengths.cumsum(0), (1, 0), value=0)

if self.guidance_scale > 1.0:
if negative_prompt is None:
negative_prompt = ""
if negative_prompt is None and negative_prompt_embeds_qwen is None:
negative_prompt = [""] * batch_size

if isinstance(negative_prompt, str):
negative_prompt = [negative_prompt] * len(prompt) if prompt is not None else [negative_prompt]
elif len(negative_prompt) != len(prompt):
negative_prompt = [negative_prompt] * batch_size
elif negative_prompt is not None and len(negative_prompt) != batch_size:
raise ValueError(
f"`negative_prompt` must have same length as `prompt`. Got {len(negative_prompt)} vs {len(prompt)}."
f"`negative_prompt` must have the same batch size as `prompt`. Got {len(negative_prompt)} vs {batch_size}."
)

if negative_prompt_embeds_qwen is None:
negative_prompt_embeds_qwen, negative_prompt_embeds_clip, negative_prompt_cu_seqlens = (
self.encode_prompt(
Expand All @@ -718,6 +734,15 @@ def __call__(
dtype=dtype,
)
)
else:
negative_prompt_embeds_qwen = negative_prompt_embeds_qwen.repeat_interleave(
num_images_per_prompt, dim=0
)
negative_prompt_embeds_clip = negative_prompt_embeds_clip.repeat_interleave(
num_images_per_prompt, dim=0
)
negative_prompt_lengths = negative_prompt_cu_seqlens.diff().repeat_interleave(num_images_per_prompt)
negative_prompt_cu_seqlens = F.pad(negative_prompt_lengths.cumsum(0), (1, 0), value=0)

# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
Expand Down
58 changes: 40 additions & 18 deletions src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -490,8 +490,6 @@ def encode_prompt(
if not isinstance(prompt, list):
prompt = [prompt]

batch_size = len(prompt)

prompt = [prompt_clean(p) for p in prompt]

# Encode with Qwen2.5-VL
Expand All @@ -513,20 +511,10 @@ def encode_prompt(

# Repeat embeddings for num_videos_per_prompt
# Qwen embeddings: repeat sequence for each video, then reshape
prompt_embeds_qwen = prompt_embeds_qwen.repeat(
1, num_videos_per_prompt, 1
) # [batch_size, seq_len * num_videos_per_prompt, embed_dim]
# Reshape to [batch_size * num_videos_per_prompt, seq_len, embed_dim]
prompt_embeds_qwen = prompt_embeds_qwen.view(
batch_size * num_videos_per_prompt, -1, prompt_embeds_qwen.shape[-1]
)
prompt_embeds_qwen = prompt_embeds_qwen.repeat_interleave(num_videos_per_prompt, dim=0)

# CLIP embeddings: repeat for each video
prompt_embeds_clip = prompt_embeds_clip.repeat(
1, num_videos_per_prompt, 1
) # [batch_size, num_videos_per_prompt, clip_embed_dim]
# Reshape to [batch_size * num_videos_per_prompt, clip_embed_dim]
prompt_embeds_clip = prompt_embeds_clip.view(batch_size * num_videos_per_prompt, -1)
prompt_embeds_clip = prompt_embeds_clip.repeat_interleave(num_videos_per_prompt, dim=0)

# Repeat cumulative sequence lengths for num_videos_per_prompt
# Original differences (lengths) for each prompt in the batch
Expand Down Expand Up @@ -602,6 +590,13 @@ def check_inputs(
"If any of `prompt_embeds_qwen`, `prompt_embeds_clip`, or `prompt_cu_seqlens` is provided, "
"all three must be provided."
)
if (
prompt_embeds_qwen.shape[0] != prompt_embeds_clip.shape[0]
or prompt_cu_seqlens.shape[0] != prompt_embeds_qwen.shape[0] + 1
):
raise ValueError(
"Precomputed positive prompt embeddings and sequence lengths must have matching batch sizes."
)

# Check for consistency within negative prompt embeddings and sequence lengths
if (
Expand All @@ -618,6 +613,13 @@ def check_inputs(
"If any of `negative_prompt_embeds_qwen`, `negative_prompt_embeds_clip`, or `negative_prompt_cu_seqlens` is provided, "
"all three must be provided."
)
if (
negative_prompt_embeds_qwen.shape[0] != negative_prompt_embeds_clip.shape[0]
or negative_prompt_cu_seqlens.shape[0] != negative_prompt_embeds_qwen.shape[0] + 1
):
raise ValueError(
"Precomputed negative prompt embeddings and sequence lengths must have matching batch sizes."
)

# Check if prompt or embeddings are provided (either prompt or all required embedding components for positive)
if prompt is None and prompt_embeds_qwen is None:
Expand Down Expand Up @@ -691,6 +693,11 @@ def prepare_latents(
# Encode the input image to use as first frame
# Preprocess image
image_tensor = self.video_processor.preprocess(image, height=height, width=width).to(device, dtype=dtype)
if batch_size % image_tensor.shape[0] != 0:
raise ValueError(
f"The effective batch size {batch_size} must be divisible by the image batch size {image_tensor.shape[0]}."
)
image_tensor = image_tensor.repeat_interleave(batch_size // image_tensor.shape[0], dim=0)

# Encode image to latents using VAE
with torch.no_grad():
Expand Down Expand Up @@ -881,17 +888,32 @@ def __call__(
device=device,
dtype=dtype,
)
else:
prompt_embeds_qwen = prompt_embeds_qwen.repeat_interleave(num_videos_per_prompt, dim=0)
prompt_embeds_clip = prompt_embeds_clip.repeat_interleave(num_videos_per_prompt, dim=0)
prompt_lengths = prompt_cu_seqlens.diff().repeat_interleave(num_videos_per_prompt)
prompt_cu_seqlens = F.pad(prompt_lengths.cumsum(0), (1, 0), value=0)

if self.guidance_scale > 1.0:
if negative_prompt is None:
if negative_prompt is None and negative_prompt_embeds_qwen is None:
negative_prompt = "Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, deformed, walking backwards"

if isinstance(negative_prompt, str):
negative_prompt = [negative_prompt] * len(prompt) if prompt is not None else [negative_prompt]
elif len(negative_prompt) != len(prompt):
negative_prompt = [negative_prompt] * batch_size
elif negative_prompt is not None and len(negative_prompt) != batch_size:
raise ValueError(
f"`negative_prompt` must have same length as `prompt`. Got {len(negative_prompt)} vs {len(prompt)}."
f"`negative_prompt` must have the same batch size as `prompt`. Got {len(negative_prompt)} vs {batch_size}."
)

if negative_prompt_embeds_qwen is not None:
negative_prompt_embeds_qwen = negative_prompt_embeds_qwen.repeat_interleave(
num_videos_per_prompt, dim=0
)
negative_prompt_embeds_clip = negative_prompt_embeds_clip.repeat_interleave(
num_videos_per_prompt, dim=0
)
negative_prompt_lengths = negative_prompt_cu_seqlens.diff().repeat_interleave(num_videos_per_prompt)
negative_prompt_cu_seqlens = F.pad(negative_prompt_lengths.cumsum(0), (1, 0), value=0)

if negative_prompt_embeds_qwen is None:
negative_prompt_embeds_qwen, negative_prompt_embeds_clip, negative_prompt_cu_seqlens = (
Expand Down
Loading
Loading