diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py index 61c9da0c9496..2984641bf689 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py @@ -446,9 +446,13 @@ def check_inputs( strength, height, width, - callback_on_step_end_tensor_inputs, + negative_prompt=None, + negative_prompt_2=None, prompt_embeds=None, + negative_prompt_embeds=None, pooled_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, max_sequence_length=None, ): if strength < 0 or strength > 1: @@ -485,10 +489,33 @@ def check_inputs( elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_2 is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_embeds is not None and pooled_prompt_embeds is None: raise ValueError( "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`." ) + if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`." + ) if max_sequence_length is not None and max_sequence_length > 512: raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}") @@ -638,6 +665,9 @@ def __call__( self, prompt: str | list[str] = None, prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, + true_cfg_scale: float = 1.0, image: PipelineImageInput = None, control_image: PipelineImageInput = None, height: int | None = None, @@ -655,6 +685,8 @@ def __call__( latents: torch.FloatTensor | None = None, prompt_embeds: torch.FloatTensor | None = None, pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: dict[str, Any] | None = None, @@ -670,6 +702,16 @@ def __call__( The prompt or prompts to guide the image generation. prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is + not greater than `1`). + negative_prompt_2 (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. + true_cfg_scale (`float`, *optional*, defaults to 1.0): + True classifier-free guidance (guidance scale) is enabled when `true_cfg_scale` > 1 and + `negative_prompt` is provided. image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): The image(s) to modify with the pipeline. control_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): @@ -711,6 +753,14 @@ def __call__( Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. pooled_prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated pooled text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generate image. Choose between `PIL.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): @@ -751,9 +801,13 @@ def __call__( strength, height, width, - callback_on_step_end_tensor_inputs, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, max_sequence_length=max_sequence_length, ) @@ -774,6 +828,10 @@ def __call__( lora_scale = ( self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None ) + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None + ) + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt ( prompt_embeds, pooled_prompt_embeds, @@ -788,6 +846,21 @@ def __call__( max_sequence_length=max_sequence_length, lora_scale=lora_scale, ) + if do_true_cfg: + ( + negative_prompt_embeds, + negative_pooled_prompt_embeds, + _, + ) = self.encode_prompt( + prompt=negative_prompt, + prompt_2=negative_prompt_2, + prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=negative_pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) init_image = self.image_processor.preprocess(image, height=height, width=width) init_image = init_image.to(dtype=torch.float32) @@ -974,6 +1047,23 @@ def __call__( controlnet_blocks_repeat=controlnet_blocks_repeat, )[0] + if do_true_cfg: + neg_noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=negative_pooled_prompt_embeds, + encoder_hidden_states=negative_prompt_embeds, + controlnet_block_samples=controlnet_block_samples, + controlnet_single_block_samples=controlnet_single_block_samples, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + controlnet_blocks_repeat=controlnet_blocks_repeat, + )[0] + noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) + latents_dtype = latents.dtype latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py index eed671152bc9..73f8e1fe00cd 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py @@ -460,8 +460,12 @@ def check_inputs( height, width, output_type, + negative_prompt=None, + negative_prompt_2=None, prompt_embeds=None, + negative_prompt_embeds=None, pooled_prompt_embeds=None, + negative_pooled_prompt_embeds=None, callback_on_step_end_tensor_inputs=None, padding_mask_crop=None, max_sequence_length=None, @@ -500,10 +504,33 @@ def check_inputs( elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_2 is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_embeds is not None and pooled_prompt_embeds is None: raise ValueError( "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`." ) + if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`." + ) if padding_mask_crop is not None: if not isinstance(image, PIL.Image.Image): @@ -742,6 +769,9 @@ def __call__( self, prompt: str | list[str] = None, prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, + true_cfg_scale: float = 1.0, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: PipelineImageInput = None, @@ -762,6 +792,8 @@ def __call__( latents: torch.FloatTensor | None = None, prompt_embeds: torch.FloatTensor | None = None, pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: dict[str, Any] | None = None, @@ -777,6 +809,16 @@ def __call__( The prompt or prompts to guide the image generation. prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is + not greater than `1`). + negative_prompt_2 (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. + true_cfg_scale (`float`, *optional*, defaults to 1.0): + True classifier-free guidance (guidance scale) is enabled when `true_cfg_scale` > 1 and + `negative_prompt` is provided. image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): The image(s) to inpaint. mask_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): @@ -825,6 +867,14 @@ def __call__( Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. pooled_prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated pooled text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generate image. Choose between `PIL.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): @@ -872,8 +922,12 @@ def __call__( height, width, output_type=output_type, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, padding_mask_crop=padding_mask_crop, max_sequence_length=max_sequence_length, @@ -898,6 +952,10 @@ def __call__( lora_scale = ( self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None ) + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None + ) + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt( prompt=prompt, prompt_2=prompt_2, @@ -908,6 +966,17 @@ def __call__( max_sequence_length=max_sequence_length, lora_scale=lora_scale, ) + if do_true_cfg: + negative_prompt_embeds, negative_pooled_prompt_embeds, _ = self.encode_prompt( + prompt=negative_prompt, + prompt_2=negative_prompt_2, + prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=negative_pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) # 4. Preprocess mask and image if padding_mask_crop is not None: @@ -1152,6 +1221,23 @@ def __call__( controlnet_blocks_repeat=controlnet_blocks_repeat, )[0] + if do_true_cfg: + neg_noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=negative_pooled_prompt_embeds, + encoder_hidden_states=negative_prompt_embeds, + controlnet_block_samples=controlnet_block_samples, + controlnet_single_block_samples=controlnet_single_block_samples, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + controlnet_blocks_repeat=controlnet_blocks_repeat, + )[0] + noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) + # compute the previous noisy sample x_t -> x_t-1 latents_dtype = latents.dtype latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py index 165e7bda2dc9..5f159172fd3e 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py @@ -158,6 +158,66 @@ def test_flux_controlnet_different_prompts(self): assert max_diff > 1e-6, "Outputs should be different for different prompts." + def test_flux_controlnet_img2img_true_cfg(self): + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + inputs.pop("generator") + no_true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + inputs["negative_prompt"] = "bad quality" + inputs["true_cfg_scale"] = 2.0 + true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + assert not torch.allclose(no_true_cfg_out, true_cfg_out), ( + "Outputs should be different when true_cfg_scale is set." + ) + + def test_flux_controlnet_img2img_true_cfg_requires_negative_prompt(self): + # `true_cfg_scale` alone must stay a no-op: true CFG only kicks in once a negative prompt + # (or precomputed negative embeddings) is supplied. + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + inputs.pop("generator") + baseline_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + inputs["true_cfg_scale"] = 2.0 + no_negative_prompt_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + assert torch.allclose(baseline_out, no_negative_prompt_out), ( + "`true_cfg_scale` should be ignored when no negative prompt is provided." + ) + + def test_flux_controlnet_img2img_true_cfg_with_negative_embeds(self): + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + inputs.pop("generator") + prompt = inputs.pop("prompt") + + prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt( + prompt=prompt, prompt_2=None, device=torch_device, num_images_per_prompt=1, max_sequence_length=48 + ) + negative_prompt_embeds, negative_pooled_prompt_embeds, _ = pipe.encode_prompt( + prompt="bad quality", prompt_2=None, device=torch_device, num_images_per_prompt=1, max_sequence_length=48 + ) + inputs.update( + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + ) + + inputs["true_cfg_scale"] = 1.0 + cfg_off = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + inputs["true_cfg_scale"] = 2.0 + cfg_on = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + assert not torch.allclose(cfg_off, cfg_on), ( + "Precomputed negative embeds should enable true CFG when negative_prompt is None." + ) + def test_fused_qkv_projections(self): # Run on CPU to keep the seeded generator deterministic across the three forward passes. pipe = self.get_pipeline() diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py index eeacc636e89d..fd56fda8c88d 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py @@ -167,6 +167,66 @@ def test_flux_controlnet_inpaint_with_controlnet_conditioning_scale(self): "Changing `controlnet_conditioning_scale` should change the output." ) + def test_flux_controlnet_inpaint_true_cfg(self): + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + inputs.pop("generator") + no_true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + inputs["negative_prompt"] = "bad quality" + inputs["true_cfg_scale"] = 2.0 + true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + assert not torch.allclose(no_true_cfg_out, true_cfg_out), ( + "Outputs should be different when true_cfg_scale is set." + ) + + def test_flux_controlnet_inpaint_true_cfg_requires_negative_prompt(self): + # `true_cfg_scale` alone must stay a no-op: true CFG only kicks in once a negative prompt + # (or precomputed negative embeddings) is supplied. + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + inputs.pop("generator") + baseline_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + inputs["true_cfg_scale"] = 2.0 + no_negative_prompt_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + assert torch.allclose(baseline_out, no_negative_prompt_out), ( + "`true_cfg_scale` should be ignored when no negative prompt is provided." + ) + + def test_flux_controlnet_inpaint_true_cfg_with_negative_embeds(self): + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + inputs.pop("generator") + prompt = inputs.pop("prompt") + + prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt( + prompt=prompt, prompt_2=None, device=torch_device, num_images_per_prompt=1, max_sequence_length=48 + ) + negative_prompt_embeds, negative_pooled_prompt_embeds, _ = pipe.encode_prompt( + prompt="bad quality", prompt_2=None, device=torch_device, num_images_per_prompt=1, max_sequence_length=48 + ) + inputs.update( + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + ) + + inputs["true_cfg_scale"] = 1.0 + cfg_off = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + inputs["true_cfg_scale"] = 2.0 + cfg_on = pipe(**inputs, generator=torch.manual_seed(0)).images[0] + + assert not torch.allclose(cfg_off, cfg_on), ( + "Precomputed negative embeds should enable true CFG when negative_prompt is None." + ) + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=3e-3): super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)