diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control.py b/src/diffusers/pipelines/flux/pipeline_flux_control.py index b50861cb437f..5bd8f402abe1 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control.py @@ -809,17 +809,18 @@ def __call__( # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] # compute the previous noisy sample x_t -> x_t-1 latents_dtype = latents.dtype diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py index 15d876caf6ca..4fda8bdcd7cc 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py @@ -892,17 +892,18 @@ def __call__( # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] # compute the previous noisy sample x_t -> x_t-1 latents_dtype = latents.dtype diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py index 6fa0c1dff9f6..8515e8c58879 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py @@ -1046,17 +1046,18 @@ def __call__( # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] # compute the previous noisy sample x_t -> x_t-1 latents_dtype = latents.dtype diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py index ab2cebb1e765..8946a2a0b3a6 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py @@ -1124,30 +1124,13 @@ def __call__( ) guidance = guidance.expand(latents.shape[0]) if guidance is not None else None - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=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] - - if do_true_cfg: - if negative_image_embeds is not None: - self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds - neg_noise_pred = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latents, timestep=timestep / 1000, guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, controlnet_block_samples=controlnet_block_samples, controlnet_single_block_samples=controlnet_single_block_samples, txt_ids=text_ids, @@ -1156,6 +1139,25 @@ def __call__( return_dict=False, controlnet_blocks_repeat=controlnet_blocks_repeat, )[0] + + if do_true_cfg: + if negative_image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds + with self.transformer.cache_context("uncond"): + 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 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 482fca3b59c3..47c62b06c323 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 @@ -961,20 +961,21 @@ def __call__( ) guidance = guidance.expand(latents.shape[0]) if guidance is not None else None - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=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] + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=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] 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 26a434807694..d8b8404c412e 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py @@ -1139,20 +1139,21 @@ def __call__( else: guidance = None - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=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] + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=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] # compute the previous noisy sample x_t -> x_t-1 latents_dtype = latents.dtype diff --git a/src/diffusers/pipelines/flux/pipeline_flux_fill.py b/src/diffusers/pipelines/flux/pipeline_flux_fill.py index 55d3464cfb4e..03726b129168 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_fill.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_fill.py @@ -958,17 +958,18 @@ def __call__( # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=torch.cat((latents, masked_image_latents), dim=2), - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=torch.cat((latents, masked_image_latents), dim=2), + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] # compute the previous noisy sample x_t -> x_t-1 latents_dtype = latents.dtype diff --git a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py index 3288bba94772..901799653dab 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py @@ -990,32 +990,34 @@ def __call__( self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - if negative_image_embeds is not None: - self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds - neg_noise_pred = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latents, timestep=timestep / 1000, guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_image_ids, joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] + + if do_true_cfg: + if negative_image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds + with self.transformer.cache_context("uncond"): + 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, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) # compute the previous noisy sample x_t -> x_t-1 diff --git a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py index 15d9bb9868fd..d268201b5f67 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py @@ -1146,32 +1146,34 @@ def __call__( self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - if negative_image_embeds is not None: - self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds - neg_noise_pred = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latents, timestep=timestep / 1000, guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_image_ids, joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] + + if do_true_cfg: + if negative_image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds + with self.transformer.cache_context("uncond"): + 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, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) # compute the previous noisy sample x_t -> x_t-1 diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py index 3579c3abd8d1..58f5170bce77 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py @@ -1038,33 +1038,35 @@ def __call__( latent_model_input = torch.cat([latents, image_latents], dim=1) timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred[:, : latents.size(1)] - - if do_true_cfg: - if negative_image_embeds is not None: - self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds - neg_noise_pred = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep / 1000, guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - txt_ids=negative_text_ids, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, img_ids=latent_ids, joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] + noise_pred = noise_pred[:, : latents.size(1)] + + if do_true_cfg: + if negative_image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds + with self.transformer.cache_context("uncond"): + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=negative_pooled_prompt_embeds, + encoder_hidden_states=negative_prompt_embeds, + txt_ids=negative_text_ids, + img_ids=latent_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] neg_noise_pred = neg_noise_pred[:, : latents.size(1)] noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py index 05bf8200f588..28c6273a73d8 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py @@ -1346,33 +1346,35 @@ def __call__( latent_model_input = torch.cat([latents, image_latents], dim=1) timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred[:, : latents.size(1)] - - if do_true_cfg: - if negative_image_embeds is not None: - self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds - neg_noise_pred = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep / 1000, guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - txt_ids=negative_text_ids, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, img_ids=latent_ids, joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] + noise_pred = noise_pred[:, : latents.size(1)] + + if do_true_cfg: + if negative_image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds + with self.transformer.cache_context("uncond"): + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=negative_pooled_prompt_embeds, + encoder_hidden_states=negative_prompt_embeds, + txt_ids=negative_text_ids, + img_ids=latent_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] neg_noise_pred = neg_noise_pred[:, : latents.size(1)] noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux.py b/tests/pipelines/controlnet_flux/test_controlnet_flux.py index 90c5e617f61b..a70a315cee4b 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux.py @@ -41,7 +41,12 @@ torch_device, ) from ..flux.testing_utils import FluxIPAdapterTesterMixin -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + TaylorSeerCacheTesterMixin, +) enable_full_determinism() @@ -55,13 +60,13 @@ class FluxControlNetPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt"]) output_shape = (3, 32, 32) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=16, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -208,6 +213,10 @@ class TestFluxControlNetPipelineMemory(FluxControlNetPipelineTesterConfig, Memor """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux ControlNet pipeline.""" +class TestFluxControlNetPipelineTaylorSeerCache(FluxControlNetPipelineTesterConfig, TaylorSeerCacheTesterMixin): + """TaylorSeer cache tests for the Flux ControlNet pipeline.""" + + @nightly @require_big_accelerator class TestFluxControlNetPipelineSlow: diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py index 165e7bda2dc9..aaff3ff333c6 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py @@ -15,6 +15,7 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + TaylorSeerCacheTesterMixin, check_qkv_fused_layers_exist, ) @@ -38,13 +39,13 @@ class FluxControlNetImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt", "image", "control_image"]) output_shape = (3, 32, 32) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=4, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -235,3 +236,9 @@ def test_flux_image_output_shape(self): class TestFluxControlNetImg2ImgPipelineMemory(FluxControlNetImg2ImgPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux ControlNet img2img pipeline.""" + + +class TestFluxControlNetImg2ImgPipelineTaylorSeerCache( + FluxControlNetImg2ImgPipelineTesterConfig, TaylorSeerCacheTesterMixin +): + """TaylorSeer cache tests for the Flux ControlNet img2img 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..d62d4a1ccb92 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py @@ -13,7 +13,12 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, floats_tensor, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + TaylorSeerCacheTesterMixin, +) enable_full_determinism() @@ -40,13 +45,13 @@ class FluxControlNetInpaintPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt", "image", "mask_image", "control_image"]) output_shape = (3, 32, 32) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=8, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -205,3 +210,9 @@ def test_flux_image_output_shape(self): class TestFluxControlNetInpaintPipelineMemory(FluxControlNetInpaintPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux ControlNet inpaint pipeline.""" + + +class TestFluxControlNetInpaintPipelineTaylorSeerCache( + FluxControlNetInpaintPipelineTesterConfig, TaylorSeerCacheTesterMixin +): + """TaylorSeer cache tests for the Flux ControlNet inpaint pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_control.py b/tests/pipelines/flux/test_pipeline_flux_control.py index 70f558b16cb2..e1cf6ff0819f 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control.py +++ b/tests/pipelines/flux/test_pipeline_flux_control.py @@ -33,6 +33,7 @@ LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, + TaylorSeerCacheTesterMixin, check_qkv_fused_layers_exist, ) @@ -45,14 +46,14 @@ class FluxControlPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt"]) output_shape = (3, 8, 8) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=8, out_channels=4, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -204,6 +205,10 @@ class TestFluxControlPipelineMemory(FluxControlPipelineTesterConfig, MemoryTeste """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux Control pipeline.""" +class TestFluxControlPipelineTaylorSeerCache(FluxControlPipelineTesterConfig, TaylorSeerCacheTesterMixin): + """TaylorSeer cache tests for the Flux Control pipeline.""" + + class TestFluxControlPipelineLoRA(FluxControlPipelineTesterConfig, LoraTesterMixin): """LoRA tests for the Flux Control pipeline. diff --git a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py index 516c6b92b504..51095ce5b2b6 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py +++ b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py @@ -10,7 +10,12 @@ ) from ...testing_utils import torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + TaylorSeerCacheTesterMixin, +) class FluxControlImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): @@ -21,14 +26,14 @@ class FluxControlImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt"]) output_shape = (3, 8, 8) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=8, out_channels=4, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -144,3 +149,9 @@ def test_flux_image_output_shape(self): class TestFluxControlImg2ImgPipelineMemory(FluxControlImg2ImgPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux control img2img pipeline.""" + + +class TestFluxControlImg2ImgPipelineTaylorSeerCache( + FluxControlImg2ImgPipelineTesterConfig, TaylorSeerCacheTesterMixin +): + """TaylorSeer cache tests for the Flux Control img2img pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py index d4e7018bc8a6..a7026d67f5b3 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py @@ -14,6 +14,7 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + TaylorSeerCacheTesterMixin, check_qkv_fused_layers_exist, ) @@ -26,14 +27,14 @@ class FluxControlInpaintPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt"]) output_shape = (3, 8, 8) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=8, out_channels=4, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -178,3 +179,9 @@ def test_flux_image_output_shape(self): class TestFluxControlInpaintPipelineMemory(FluxControlInpaintPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux control inpaint pipeline.""" + + +class TestFluxControlInpaintPipelineTaylorSeerCache( + FluxControlInpaintPipelineTesterConfig, TaylorSeerCacheTesterMixin +): + """TaylorSeer cache tests for the Flux Control inpaint pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_fill.py b/tests/pipelines/flux/test_pipeline_flux_fill.py index 4dc3aa5f51d6..d12d0341e60d 100644 --- a/tests/pipelines/flux/test_pipeline_flux_fill.py +++ b/tests/pipelines/flux/test_pipeline_flux_fill.py @@ -6,7 +6,12 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxFillPipeline, FluxTransformer2DModel from ...testing_utils import floats_tensor, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + TaylorSeerCacheTesterMixin, +) class FluxFillPipelineTesterConfig(BasePipelineTesterConfig): @@ -17,14 +22,14 @@ class FluxFillPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt"]) output_shape = (3, 32, 32) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=20, out_channels=8, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -142,3 +147,7 @@ def test_inference_batch_single_identical(self): class TestFluxFillPipelineMemory(FluxFillPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux fill pipeline.""" + + +class TestFluxFillPipelineTaylorSeerCache(FluxFillPipelineTesterConfig, TaylorSeerCacheTesterMixin): + """TaylorSeer cache tests for the Flux Fill pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_img2img.py b/tests/pipelines/flux/test_pipeline_flux_img2img.py index 93736df73ff5..b33e450571e4 100644 --- a/tests/pipelines/flux/test_pipeline_flux_img2img.py +++ b/tests/pipelines/flux/test_pipeline_flux_img2img.py @@ -6,7 +6,12 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxImg2ImgPipeline, FluxTransformer2DModel from ...testing_utils import floats_tensor, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + TaylorSeerCacheTesterMixin, +) from .testing_utils import FluxIPAdapterTesterMixin @@ -18,13 +23,13 @@ class FluxImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt"]) output_shape = (3, 8, 8) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=4, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -171,3 +176,7 @@ class TestFluxImg2ImgPipelineIPAdapter(FluxImg2ImgPipelineTesterConfig, FluxIPAd class TestFluxImg2ImgPipelineMemory(FluxImg2ImgPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux img2img pipeline.""" + + +class TestFluxImg2ImgPipelineTaylorSeerCache(FluxImg2ImgPipelineTesterConfig, TaylorSeerCacheTesterMixin): + """TaylorSeer cache tests for the Flux img2img pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_inpaint.py index 87e778083134..3f5a7bcc2edd 100644 --- a/tests/pipelines/flux/test_pipeline_flux_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_inpaint.py @@ -6,7 +6,12 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxInpaintPipeline, FluxTransformer2DModel from ...testing_utils import floats_tensor, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + TaylorSeerCacheTesterMixin, +) from .testing_utils import FluxIPAdapterTesterMixin @@ -18,13 +23,13 @@ class FluxInpaintPipelineTesterConfig(BasePipelineTesterConfig): batch_input_params = frozenset(["prompt"]) output_shape = (3, 32, 32) - def get_dummy_components(self): + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = FluxTransformer2DModel( patch_size=1, in_channels=8, - num_layers=1, - num_single_layers=1, + num_layers=num_layers, + num_single_layers=num_single_layers, attention_head_dim=16, num_attention_heads=2, joint_attention_dim=32, @@ -146,3 +151,7 @@ class TestFluxInpaintPipelineIPAdapter(FluxInpaintPipelineTesterConfig, FluxIPAd class TestFluxInpaintPipelineMemory(FluxInpaintPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux inpaint pipeline.""" + + +class TestFluxInpaintPipelineTaylorSeerCache(FluxInpaintPipelineTesterConfig, TaylorSeerCacheTesterMixin): + """TaylorSeer cache tests for the Flux inpaint pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext.py b/tests/pipelines/flux/test_pipeline_flux_kontext.py index 1fb5a52f7c84..bff2b62c7089 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext.py @@ -16,6 +16,7 @@ MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, + TaylorSeerCacheTesterMixin, ) from .testing_utils import FluxIPAdapterTesterMixin @@ -186,3 +187,7 @@ class TestFluxKontextPipelineFasterCache(FluxKontextPipelineTesterConfig, Faster "attention_weight_callback": lambda _: 0.5, "is_guidance_distilled": True, } + + +class TestFluxKontextPipelineTaylorSeerCache(FluxKontextPipelineTesterConfig, TaylorSeerCacheTesterMixin): + """TaylorSeer cache tests for the Flux Kontext pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py index cfae5b4a2a38..7743c3280279 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py @@ -17,6 +17,7 @@ MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, + TaylorSeerCacheTesterMixin, ) from .testing_utils import FluxIPAdapterTesterMixin @@ -184,6 +185,12 @@ class TestFluxKontextInpaintPipelineMemory(FluxKontextInpaintPipelineTesterConfi """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux Kontext inpaint pipeline.""" +class TestFluxKontextInpaintPipelineTaylorSeerCache( + FluxKontextInpaintPipelineTesterConfig, TaylorSeerCacheTesterMixin +): + """TaylorSeer cache tests for the Flux Kontext inpaint pipeline.""" + + class TestFluxKontextInpaintPipelinePyramidAttentionBroadcast( FluxKontextInpaintPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin ):