From a888d4904a75184cd64416ad079bb06b4ae4f546 Mon Sep 17 00:00:00 2001 From: Pawel Date: Fri, 10 Jul 2026 18:00:02 -0700 Subject: [PATCH] Fix Kandinsky 5 conditioning batch support --- .../kandinsky5/pipeline_kandinsky.py | 54 +++++++++----- .../kandinsky5/pipeline_kandinsky_i2i.py | 71 +++++++++++++------ .../kandinsky5/pipeline_kandinsky_i2v.py | 58 ++++++++++----- .../kandinsky5/pipeline_kandinsky_t2i.py | 54 +++++++++----- tests/pipelines/kandinsky5/test_kandinsky5.py | 29 ++++++++ .../kandinsky5/test_kandinsky5_i2i.py | 42 ++++++++++- .../kandinsky5/test_kandinsky5_i2v.py | 29 ++++++++ .../kandinsky5/test_kandinsky5_t2i.py | 51 +++++++++++++ 8 files changed, 308 insertions(+), 80 deletions(-) diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py index 1ce885b21f5b..83ed9b480424 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py @@ -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 @@ -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, ...] @@ -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 ( @@ -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: @@ -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) diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py index 0a8382d6031f..64d280fbd1b2 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py @@ -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 @@ -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 @@ -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 @@ -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 ( @@ -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: @@ -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) @@ -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, @@ -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( @@ -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) diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py index e82dc737f1a9..be43ab898079 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py @@ -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 @@ -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 @@ -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 ( @@ -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: @@ -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(): @@ -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 = ( diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py index 2a58d4bed33a..6bf9359399cd 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py @@ -325,8 +325,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 @@ -348,20 +346,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 @@ -435,6 +423,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 ( @@ -451,6 +446,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: @@ -654,16 +656,21 @@ 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: @@ -676,6 +683,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) diff --git a/tests/pipelines/kandinsky5/test_kandinsky5.py b/tests/pipelines/kandinsky5/test_kandinsky5.py index 4101e7798dea..9ca7faf9f1d3 100644 --- a/tests/pipelines/kandinsky5/test_kandinsky5.py +++ b/tests/pipelines/kandinsky5/test_kandinsky5.py @@ -194,6 +194,35 @@ def test_inference(self): self.assertEqual(video.shape, (3, 3, 16, 16)) + def test_num_videos_per_prompt(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + + frames = pipe(**inputs, num_videos_per_prompt=2).frames + + self.assertEqual(frames.shape[0], 2) + + def test_precomputed_embeddings_with_classifier_free_guidance(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + prompt_embeds = pipe.encode_prompt(inputs.pop("prompt"), max_sequence_length=8) + negative_prompt_embeds = pipe.encode_prompt("", max_sequence_length=8) + inputs.update( + prompt_embeds_qwen=prompt_embeds[0], + prompt_embeds_clip=prompt_embeds[1], + prompt_cu_seqlens=prompt_embeds[2], + negative_prompt_embeds_qwen=negative_prompt_embeds[0], + negative_prompt_embeds_clip=negative_prompt_embeds[1], + negative_prompt_cu_seqlens=negative_prompt_embeds[2], + ) + inputs["num_videos_per_prompt"] = 2 + + frames = pipe(**inputs).frames + + self.assertEqual(frames.shape[0], 2) + def test_attention_slicing_forward_pass(self): pass diff --git a/tests/pipelines/kandinsky5/test_kandinsky5_i2i.py b/tests/pipelines/kandinsky5/test_kandinsky5_i2i.py index dc832990836f..5a2c746d1dda 100644 --- a/tests/pipelines/kandinsky5/test_kandinsky5_i2i.py +++ b/tests/pipelines/kandinsky5/test_kandinsky5_i2i.py @@ -196,9 +196,47 @@ def test_inference(self): def test_encode_prompt_works_in_isolation(self): pass - @unittest.skip("TODO: revisit, Batch isnot yet supported in this pipeline") def test_num_images_per_prompt(self): - pass + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.resolutions = [(64, 64)] + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + + image = pipe(**inputs, num_images_per_prompt=2).image + + self.assertEqual(image.shape, (2, 3, 64, 64)) + + def test_precomputed_embeddings_with_classifier_free_guidance(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.resolutions = [(64, 64)] + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + prompt_embeds = pipe.encode_prompt(inputs.pop("prompt"), image=inputs["image"], max_sequence_length=8) + negative_prompt_embeds = pipe.encode_prompt("", image=inputs["image"], max_sequence_length=8) + inputs.update( + prompt_embeds_qwen=prompt_embeds[0], + prompt_embeds_clip=prompt_embeds[1], + prompt_cu_seqlens=prompt_embeds[2], + negative_prompt_embeds_qwen=negative_prompt_embeds[0], + negative_prompt_embeds_clip=negative_prompt_embeds[1], + negative_prompt_cu_seqlens=negative_prompt_embeds[2], + ) + inputs["num_images_per_prompt"] = 2 + + image = pipe(**inputs).image + + self.assertEqual(image.shape, (2, 3, 64, 64)) + + def test_tensor_image_input(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.resolutions = [(64, 64)] + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + inputs["image"] = torch.zeros(1, 3, 64, 64) + + image = pipe(**inputs).image + + self.assertEqual(image.shape, (1, 3, 64, 64)) @unittest.skip("TODO: revisit, Batch isnot yet supported in this pipeline") def test_inference_batch_single_identical(self): diff --git a/tests/pipelines/kandinsky5/test_kandinsky5_i2v.py b/tests/pipelines/kandinsky5/test_kandinsky5_i2v.py index 483c7b66e07b..785c970e3f50 100644 --- a/tests/pipelines/kandinsky5/test_kandinsky5_i2v.py +++ b/tests/pipelines/kandinsky5/test_kandinsky5_i2v.py @@ -198,6 +198,35 @@ def test_inference(self): # 17 frames, RGB, 32×32 self.assertEqual(video.shape, (17, 3, 32, 32)) + def test_num_videos_per_prompt(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + + frames = pipe(**inputs, num_videos_per_prompt=2).frames + + self.assertEqual(frames.shape[0], 2) + + def test_precomputed_embeddings_with_classifier_free_guidance(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + prompt_embeds = pipe.encode_prompt(inputs.pop("prompt"), max_sequence_length=8) + negative_prompt_embeds = pipe.encode_prompt("", max_sequence_length=8) + inputs.update( + prompt_embeds_qwen=prompt_embeds[0], + prompt_embeds_clip=prompt_embeds[1], + prompt_cu_seqlens=prompt_embeds[2], + negative_prompt_embeds_qwen=negative_prompt_embeds[0], + negative_prompt_embeds_clip=negative_prompt_embeds[1], + negative_prompt_cu_seqlens=negative_prompt_embeds[2], + ) + inputs["num_videos_per_prompt"] = 2 + + frames = pipe(**inputs).frames + + self.assertEqual(frames.shape[0], 2) + @unittest.skip("TODO:Test does not work") def test_encode_prompt_works_in_isolation(self): pass diff --git a/tests/pipelines/kandinsky5/test_kandinsky5_t2i.py b/tests/pipelines/kandinsky5/test_kandinsky5_t2i.py index e961103906a2..f48e81719a06 100644 --- a/tests/pipelines/kandinsky5/test_kandinsky5_t2i.py +++ b/tests/pipelines/kandinsky5/test_kandinsky5_t2i.py @@ -190,6 +190,57 @@ def test_inference(self): def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=5e-3) + def test_num_images_per_prompt(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.resolutions = [(64, 64)] + pipe.set_progress_bar_config(disable=None) + inputs = self.get_dummy_inputs("cpu") + inputs["num_images_per_prompt"] = 2 + + image = pipe(**inputs).image + + self.assertEqual(image.shape, (2, 3, 16, 16)) + + def test_precomputed_embeddings_with_classifier_free_guidance(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.resolutions = [(64, 64)] + pipe.set_progress_bar_config(disable=None) + prompt_embeds = pipe.encode_prompt("a red square", max_sequence_length=8) + negative_prompt_embeds = pipe.encode_prompt("", max_sequence_length=8) + inputs = self.get_dummy_inputs("cpu") + inputs.pop("prompt") + inputs.update( + prompt_embeds_qwen=prompt_embeds[0], + prompt_embeds_clip=prompt_embeds[1], + prompt_cu_seqlens=prompt_embeds[2], + negative_prompt_embeds_qwen=negative_prompt_embeds[0], + negative_prompt_embeds_clip=negative_prompt_embeds[1], + negative_prompt_cu_seqlens=negative_prompt_embeds[2], + num_images_per_prompt=2, + ) + + image = pipe(**inputs).image + + self.assertEqual(image.shape, (2, 3, 16, 16)) + + def test_precomputed_embeddings_with_default_negative_prompt(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to("cpu") + pipe.resolutions = [(64, 64)] + pipe.set_progress_bar_config(disable=None) + prompt_embeds = pipe.encode_prompt(["a red square", "a blue circle"], max_sequence_length=8) + inputs = self.get_dummy_inputs("cpu") + inputs.pop("prompt") + inputs.update( + prompt_embeds_qwen=prompt_embeds[0], + prompt_embeds_clip=prompt_embeds[1], + prompt_cu_seqlens=prompt_embeds[2], + num_images_per_prompt=2, + ) + + image = pipe(**inputs).image + + self.assertEqual(image.shape, (4, 3, 16, 16)) + @unittest.skip("Test not supported") def test_attention_slicing_forward_pass(self): pass