From 43e6c2087a6f07f41289aa792a7aa46bbcbbedc3 Mon Sep 17 00:00:00 2001 From: HyoungwonCho Date: Tue, 9 Jul 2024 01:20:12 +0900 Subject: [PATCH 01/27] original sd3 pipelien --- .../pipeline_stable_diffusion_3 copy.py | 904 ++++++++++++++++++ 1 file changed, 904 insertions(+) create mode 100644 src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3 copy.py diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3 copy.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3 copy.py new file mode 100644 index 000000000000..3980d323c94c --- /dev/null +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3 copy.py @@ -0,0 +1,904 @@ +# Copyright 2024 Stability AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +from transformers import ( + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, + T5TokenizerFast, +) + +from ...image_processor import VaeImageProcessor +from ...loaders import FromSingleFileMixin, SD3LoraLoaderMixin +from ...models.autoencoders import AutoencoderKL +from ...models.transformers import SD3Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import ( + is_torch_xla_available, + logging, + replace_example_docstring, +) +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import StableDiffusion3PipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import StableDiffusion3Pipeline + + >>> pipe = StableDiffusion3Pipeline.from_pretrained( + ... "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16 + ... ) + >>> pipe.to("cuda") + >>> prompt = "A cat holding a sign that says hello world" + >>> image = pipe(prompt).images[0] + >>> image.save("sd3.png") + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin): + r""" + Args: + transformer ([`SD3Transformer2DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`CLIPTextModelWithProjection`]): + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), + specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, + with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` + as its dimension. + text_encoder_2 ([`CLIPTextModelWithProjection`]): + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), + specifically the + [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k) + variant. + text_encoder_3 ([`T5EncoderModel`]): + Frozen text-encoder. Stable Diffusion 3 uses + [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the + [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant. + tokenizer (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). + tokenizer_2 (`CLIPTokenizer`): + Second Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). + tokenizer_3 (`T5TokenizerFast`): + Tokenizer of class + [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->transformer->vae" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] + + def __init__( + self, + transformer: SD3Transformer2DModel, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: CLIPTextModelWithProjection, + tokenizer: CLIPTokenizer, + text_encoder_2: CLIPTextModelWithProjection, + tokenizer_2: CLIPTokenizer, + text_encoder_3: T5EncoderModel, + tokenizer_3: T5TokenizerFast, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + text_encoder_3=text_encoder_3, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + tokenizer_3=tokenizer_3, + transformer=transformer, + scheduler=scheduler, + ) + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 + ) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.tokenizer_max_length = ( + self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 + ) + self.default_sample_size = ( + self.transformer.config.sample_size + if hasattr(self, "transformer") and self.transformer is not None + else 128 + ) + + def _get_t5_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + num_images_per_prompt: int = 1, + max_sequence_length: int = 256, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if self.text_encoder_3 is None: + return torch.zeros( + ( + batch_size * num_images_per_prompt, + self.tokenizer_max_length, + self.transformer.config.joint_attention_dim, + ), + device=device, + dtype=dtype, + ) + + text_inputs = self.tokenizer_3( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {max_sequence_length} tokens: {removed_text}" + ) + + prompt_embeds = self.text_encoder_3(text_input_ids.to(device))[0] + + dtype = self.text_encoder_3.dtype + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + _, seq_len, _ = prompt_embeds.shape + + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + return prompt_embeds + + def _get_clip_prompt_embeds( + self, + prompt: Union[str, List[str]], + num_images_per_prompt: int = 1, + device: Optional[torch.device] = None, + clip_skip: Optional[int] = None, + clip_model_index: int = 0, + ): + device = device or self._execution_device + + clip_tokenizers = [self.tokenizer, self.tokenizer_2] + clip_text_encoders = [self.text_encoder, self.text_encoder_2] + + tokenizer = clip_tokenizers[clip_model_index] + text_encoder = clip_text_encoders[clip_model_index] + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer_max_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer_max_length} tokens: {removed_text}" + ) + prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) + pooled_prompt_embeds = prompt_embeds[0] + + if clip_skip is None: + prompt_embeds = prompt_embeds.hidden_states[-2] + else: + prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] + + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) + + _, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + return prompt_embeds, pooled_prompt_embeds + + def encode_prompt( + self, + prompt: Union[str, List[str]], + prompt_2: Union[str, List[str]], + prompt_3: Union[str, List[str]], + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + negative_prompt_3: Optional[Union[str, List[str]]] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + clip_skip: Optional[int] = None, + max_sequence_length: int = 256, + ): + r""" + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + used in all text-encoders + prompt_3 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is + used in all text-encoders + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + 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 `guidance_scale` is + less 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. + negative_prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and + `text_encoder_3`. If not defined, `negative_prompt` is used in both text-encoders + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + 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. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `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. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_2 = prompt_2 or prompt + prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 + + prompt_3 = prompt_3 or prompt + prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 + + prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + clip_skip=clip_skip, + clip_model_index=0, + ) + prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( + prompt=prompt_2, + device=device, + num_images_per_prompt=num_images_per_prompt, + clip_skip=clip_skip, + clip_model_index=1, + ) + clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) + + t5_prompt_embed = self._get_t5_prompt_embeds( + prompt=prompt_3, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + clip_prompt_embeds = torch.nn.functional.pad( + clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) + ) + + prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) + pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt_2 = negative_prompt_2 or negative_prompt + negative_prompt_3 = negative_prompt_3 or negative_prompt + + # normalize str to list + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_2 = ( + batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 + ) + negative_prompt_3 = ( + batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 + ) + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( + negative_prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + clip_skip=None, + clip_model_index=0, + ) + negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( + negative_prompt_2, + device=device, + num_images_per_prompt=num_images_per_prompt, + clip_skip=None, + clip_model_index=1, + ) + negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) + + t5_negative_prompt_embed = self._get_t5_prompt_embeds( + prompt=negative_prompt_3, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + negative_clip_prompt_embeds = torch.nn.functional.pad( + negative_clip_prompt_embeds, + (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), + ) + + negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) + negative_pooled_prompt_embeds = torch.cat( + [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 + ) + + return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds + + def check_inputs( + self, + prompt, + prompt_2, + prompt_3, + height, + width, + negative_prompt=None, + negative_prompt_2=None, + negative_prompt_3=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 height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt_2 is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt_3 is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_3`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + 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)}") + elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): + raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") + + 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." + ) + elif negative_prompt_3 is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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}") + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + if latents is not None: + return latents.to(device=device, dtype=dtype) + + shape = ( + batch_size, + num_channels_latents, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def clip_skip(self): + return self._clip_skip + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + prompt_3: Optional[Union[str, List[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 28, + timesteps: List[int] = None, + guidance_scale: float = 7.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + negative_prompt_3: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 256, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + will be used instead + prompt_3 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is + will be used instead + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + 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 `guidance_scale` is + less 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 instead + negative_prompt_3 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and + `text_encoder_3`. If not defined, `negative_prompt` is used instead + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + 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. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `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](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead + of a plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`: + [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] if `return_dict` is True, otherwise a + `tuple`. When returning a tuple, the first element is a list with the generated images. + """ + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + prompt_3, + height, + width, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + 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, + ) + + self._guidance_scale = guidance_scale + self._clip_skip = clip_skip + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + ( + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + prompt_3=prompt_3, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + device=device, + clip_skip=self.clip_skip, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + pooled_projections=pooled_prompt_embeds, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + # 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] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + negative_pooled_prompt_embeds = callback_outputs.pop( + "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds + ) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if output_type == "latent": + image = latents + + else: + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return StableDiffusion3PipelineOutput(images=image) From 24394c59bae3d83d52b762c4fe9df75c3216eeae Mon Sep 17 00:00:00 2001 From: sunovivid Date: Sun, 14 Jul 2024 03:10:31 +0900 Subject: [PATCH 02/27] Add PAG support sd3 --- src/diffusers/models/attention_processor.py | 294 ++++++++++++++++++ src/diffusers/pipelines/auto_pipeline.py | 2 + src/diffusers/pipelines/pag/__init__.py | 2 + src/diffusers/pipelines/pag/pag_utils.py | 78 +++++ .../pipeline_pag_sd_3.py} | 46 ++- 5 files changed, 414 insertions(+), 8 deletions(-) rename src/diffusers/pipelines/{stable_diffusion_3/pipeline_stable_diffusion_3 copy.py => pag/pipeline_pag_sd_3.py} (95%) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index ac773ba48103..f9d270dcd7d6 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1058,6 +1058,300 @@ def __call__( return hidden_states, encoder_hidden_states +class PAGJointAttnProcessor2_0: + """Attention processor used typically in processing the SD3-like self-attention projections.""" + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + *args, + **kwargs, + ) -> torch.FloatTensor: + residual = hidden_states + + input_ndim = hidden_states.ndim + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + context_input_ndim = encoder_hidden_states.ndim + if context_input_ndim == 4: + batch_size, channel, height, width = encoder_hidden_states.shape + encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + identity_block_size = hidden_states.shape[1] # patch embeddings width * height (correspond to self-attention map width or height) + + # chunk + hidden_states_org, hidden_states_ptb = hidden_states.chunk(2) + encoder_hidden_states_org, encoder_hidden_states_ptb = encoder_hidden_states.chunk(2) + + ################## original path ################## + batch_size = encoder_hidden_states_org.shape[0] + + # `sample` projections. + query = attn.to_q(hidden_states_org) + key = attn.to_k(hidden_states_org) + value = attn.to_v(hidden_states_org) + + # `context` projections. + encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org) + encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org) + encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org) + + # attention + query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) + key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) + value = torch.cat([value, encoder_hidden_states_org_value_proj], dim=1) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + hidden_states_org = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states_org = hidden_states_org.to(query.dtype) + + # Split the attention outputs. + hidden_states_org, encoder_hidden_states_org = ( + hidden_states_org[:, : residual.shape[1]], + hidden_states_org[:, residual.shape[1] :], + ) + + # linear proj + hidden_states_org = attn.to_out[0](hidden_states_org) + # dropout + hidden_states_org = attn.to_out[1](hidden_states_org) + if not attn.context_pre_only: + encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org) + + if input_ndim == 4: + hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) + if context_input_ndim == 4: + encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) + + ################## perturbed path ################## + + batch_size = encoder_hidden_states_ptb.shape[0] + + # `sample` projections. + query = attn.to_q(hidden_states_ptb) + key = attn.to_k(hidden_states_ptb) + value = attn.to_v(hidden_states_ptb) + + # `context` projections. + encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) + encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) + encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb) + + # attention + query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) + key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) + value = torch.cat([value, encoder_hidden_states_ptb_value_proj], dim=1) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # create a full mask with all entries set to 0 + seq_len = query.size(2) + full_mask = torch.zeros((seq_len, seq_len), device=query.device, dtype=query.dtype) + + # set the attention value between image patches to -inf + full_mask[:identity_block_size, :identity_block_size] = float('-inf') + + # set the diagonal of the attention value between image patches to 0 + full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0) + + # expand the mask to match the attention weights shape + full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions + + hidden_states_ptb = F.scaled_dot_product_attention(query, key, value, attn_mask=full_mask, dropout_p=0.0, is_causal=False) + hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states_ptb = hidden_states_ptb.to(query.dtype) + + # split the attention outputs. + hidden_states_ptb, encoder_hidden_states_ptb = ( + hidden_states_ptb[:, : residual.shape[1]], + hidden_states_ptb[:, residual.shape[1] :], + ) + + # linear proj + hidden_states_ptb = attn.to_out[0](hidden_states_ptb) + # dropout + hidden_states_ptb = attn.to_out[1](hidden_states_ptb) + if not attn.context_pre_only: + encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb) + + if input_ndim == 4: + hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) + if context_input_ndim == 4: + encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) + + ################ concat ############### + hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) + encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb]) + + return hidden_states, encoder_hidden_states + + +class PAGCFGJointAttnProcessor2_0: + """Attention processor used typically in processing the SD3-like self-attention projections.""" + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + attention_mask: Optional[torch.FloatTensor] = None, + *args, + **kwargs, + ) -> torch.FloatTensor: + residual = hidden_states + + input_ndim = hidden_states.ndim + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + context_input_ndim = encoder_hidden_states.ndim + if context_input_ndim == 4: + batch_size, channel, height, width = encoder_hidden_states.shape + encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + identity_block_size = hidden_states.shape[1] # patch embeddings width * height (correspond to self-attention map width or height) + + # chunk + hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3) + hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org]) + + encoder_hidden_states_uncond, encoder_hidden_states_org, encoder_hidden_states_ptb =encoder_hidden_states.chunk(3) + encoder_hidden_states_org = torch.cat([encoder_hidden_states_uncond, encoder_hidden_states_org]) + + ################## original path ################## + batch_size = encoder_hidden_states_org.shape[0] + + # `sample` projections. + query = attn.to_q(hidden_states_org) + key = attn.to_k(hidden_states_org) + value = attn.to_v(hidden_states_org) + + # `context` projections. + encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org) + encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org) + encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org) + + # attention + query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) + key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) + value = torch.cat([value, encoder_hidden_states_org_value_proj], dim=1) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + hidden_states_org = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states_org = hidden_states_org.to(query.dtype) + + # Split the attention outputs. + hidden_states_org, encoder_hidden_states_org = ( + hidden_states_org[:, : residual.shape[1]], + hidden_states_org[:, residual.shape[1] :], + ) + + # linear proj + hidden_states_org = attn.to_out[0](hidden_states_org) + # dropout + hidden_states_org = attn.to_out[1](hidden_states_org) + if not attn.context_pre_only: + encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org) + + if input_ndim == 4: + hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) + if context_input_ndim == 4: + encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) + + ################## perturbed path ################## + + batch_size = encoder_hidden_states_ptb.shape[0] + + # `sample` projections. + query = attn.to_q(hidden_states_ptb) + key = attn.to_k(hidden_states_ptb) + value = attn.to_v(hidden_states_ptb) + + # `context` projections. + encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) + encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) + encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb) + + # attention + query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) + key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) + value = torch.cat([value, encoder_hidden_states_ptb_value_proj], dim=1) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # create a full mask with all entries set to 0 + seq_len = query.size(2) + full_mask = torch.zeros((seq_len, seq_len), device=query.device, dtype=query.dtype) + + # set the attention value between image patches to -inf + full_mask[:identity_block_size, :identity_block_size] = float('-inf') + + # set the diagonal of the attention value between image patches to 0 + full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0) + + # expand the mask to match the attention weights shape + full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions + + hidden_states_ptb = F.scaled_dot_product_attention(query, key, value, attn_mask=full_mask, dropout_p=0.0, is_causal=False) + hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states_ptb = hidden_states_ptb.to(query.dtype) + + # split the attention outputs. + hidden_states_ptb, encoder_hidden_states_ptb = ( + hidden_states_ptb[:, : residual.shape[1]], + hidden_states_ptb[:, residual.shape[1] :], + ) + + # linear proj + hidden_states_ptb = attn.to_out[0](hidden_states_ptb) + # dropout + hidden_states_ptb = attn.to_out[1](hidden_states_ptb) + if not attn.context_pre_only: + encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb) + + if input_ndim == 4: + hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) + if context_input_ndim == 4: + encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) + + ################ concat ############### + hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) + encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb]) + + return hidden_states, encoder_hidden_states + + class FusedJointAttnProcessor2_0: """Attention processor used typically in processing the SD3-like self-attention projections.""" diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index 3eb98cfef912..a312eed89559 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -52,6 +52,7 @@ StableDiffusionXLPAGImg2ImgPipeline, StableDiffusionXLPAGInpaintPipeline, StableDiffusionXLPAGPipeline, + StableDiffusion3PAGPipeline, ) from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline from .stable_cascade import StableCascadeCombinedPipeline, StableCascadeDecoderPipeline @@ -92,6 +93,7 @@ ("stable-diffusion-pag", StableDiffusionPAGPipeline), ("stable-diffusion-xl-pag", StableDiffusionXLPAGPipeline), ("stable-diffusion-xl-controlnet-pag", StableDiffusionXLControlNetPAGPipeline), + ("stable-diffusion-3-medium-diffusers", StableDiffusion3PAGPipeline) ] ) diff --git a/src/diffusers/pipelines/pag/__init__.py b/src/diffusers/pipelines/pag/__init__.py index 5989a6237d40..1a6b50f585d8 100644 --- a/src/diffusers/pipelines/pag/__init__.py +++ b/src/diffusers/pipelines/pag/__init__.py @@ -27,6 +27,7 @@ _import_structure["pipeline_pag_sd_xl"] = ["StableDiffusionXLPAGPipeline"] _import_structure["pipeline_pag_sd_xl_img2img"] = ["StableDiffusionXLPAGImg2ImgPipeline"] _import_structure["pipeline_pag_sd_xl_inpaint"] = ["StableDiffusionXLPAGInpaintPipeline"] + _import_structure["pipeline_pag_sd_3"] = ["StableDiffusion3PAGPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -41,6 +42,7 @@ from .pipeline_pag_sd_xl import StableDiffusionXLPAGPipeline from .pipeline_pag_sd_xl_img2img import StableDiffusionXLPAGImg2ImgPipeline from .pipeline_pag_sd_xl_inpaint import StableDiffusionXLPAGInpaintPipeline + from .pipeline_pag_sd_3 import StableDiffusion3PAGPipeline else: import sys diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py index 2009024e4e47..bf5e2fb87e0e 100644 --- a/src/diffusers/pipelines/pag/pag_utils.py +++ b/src/diffusers/pipelines/pag/pag_utils.py @@ -17,6 +17,8 @@ from ...models.attention_processor import ( PAGCFGIdentitySelfAttnProcessor2_0, PAGIdentitySelfAttnProcessor2_0, + PAGJointAttnProcessor2_0, + PAGCFGJointAttnProcessor2_0, ) from ...utils import logging @@ -55,6 +57,24 @@ def _check_input_pag_applied_layer(layer): if not layer_splits[2].startswith("attentions_"): raise ValueError(f"Invalid attention_index in pag layer: {layer}. Should start with 'attentions_'") + @staticmethod + def _check_input_pag_applied_layer_sd3(layer): + r""" + Check if each layer input in `applied_pag_layers` is valid. It should be the block index: {block_index}. + """ + + # Check if the layer index is valid (should be int or str of int) + if isinstance(layer, int): + return # Valid layer index + + if isinstance(layer, str): + if layer.isdigit(): + return # Valid layer index + + # If it is not a valid layer index, raise a ValueError + raise ValueError(f"Pag layer should only contain block index. Accept number string like '3', got {layer}") + + def _set_pag_attn_processor(self, pag_applied_layers, do_classifier_free_guidance): r""" Set the attention processor for the PAG layers. @@ -146,6 +166,51 @@ def get_attn_index(module_name): for module in target_modules: module.processor = pag_attn_proc + + def _set_pag_attn_processor_sd3(self, pag_applied_layers, do_classifier_free_guidance): + r""" + Set the attention processor for the PAG layers. + """ + if do_classifier_free_guidance: + pag_attn_proc = PAGCFGJointAttnProcessor2_0() + else: + pag_attn_proc = PAGJointAttnProcessor2_0() + + def is_attn(module_name): + r""" + Check if the module is self-attention module based on its name. + """ + return "attn" in module_name and len(module_name.split(".")) == 3 # include transformer_blocks.1.attn, exclude transformer_blocks.18.attn.to_q, transformer_blocks.1.attn.add_q_proj, ... + + def get_block_index(module_name): + r""" + Get the block index from the module name. can be "block_0", "block_1", ... If there is only one block (e.g. + mid_block) and index is ommited from the name, it will be "block_0". + """ + # transformer_blocks.23.attn -> "23" + return module_name.split('.')[1] + + for pag_layer_input in pag_applied_layers: + # for each PAG layer input, we find corresponding self-attention layers in the transformer model + target_modules = [] + + block_index = str(pag_layer_input) + + for name, module in self.transformer.named_modules(): + if ( + is_attn(name) + and get_block_index(name) == block_index + ): + target_modules.append(module) + + + if len(target_modules) == 0: + raise ValueError(f"Cannot find pag layer to set attention processor for: {pag_layer_input}") + + for module in target_modules: + module.processor = pag_attn_proc + + def _get_pag_scale(self, t): r""" Get the scale factor for the perturbed attention guidance at timestep `t`. @@ -216,6 +281,19 @@ def set_pag_applied_layers(self, pag_applied_layers): self._check_input_pag_applied_layer(pag_layer) self.pag_applied_layers = pag_applied_layers + + def set_pag_applied_layers_sd3(self, pag_applied_layers): + r""" + set the the self-attention layers to apply PAG. Raise ValueError if the input is invalid. + """ + + if not isinstance(pag_applied_layers, list): + pag_applied_layers = [pag_applied_layers] + + for pag_layer in pag_applied_layers: + self._check_input_pag_applied_layer_sd3(pag_layer) + + self.pag_applied_layers = pag_applied_layers @property def pag_scale(self): diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3 copy.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py similarity index 95% rename from src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3 copy.py rename to src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index 3980d323c94c..bb72fc1f1de8 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3 copy.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -35,7 +35,8 @@ ) from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import StableDiffusion3PipelineOutput +from ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput +from .pag_utils import PAGMixin if is_torch_xla_available(): @@ -125,7 +126,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin): +class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, PAGMixin): r""" Args: transformer ([`SD3Transformer2DModel`]): @@ -174,8 +175,8 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, + pag_applied_layers: Union[str, List[str]] = "11", # ["11"] ): - super().__init__() self.register_modules( vae=vae, @@ -200,6 +201,9 @@ def __init__( if hasattr(self, "transformer") and self.transformer is not None else 128 ) + + self.set_pag_applied_layers_sd3(pag_applied_layers) + def _get_t5_prompt_embeds( self, @@ -658,6 +662,8 @@ def __call__( callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = ["latents"], max_sequence_length: int = 256, + pag_scale: float = 3.0, + pag_adaptive_scale: float = 0.0, ): r""" Function invoked when calling the pipeline for generation. @@ -776,7 +782,9 @@ def __call__( self._clip_skip = clip_skip self._joint_attention_kwargs = joint_attention_kwargs self._interrupt = False - + self._pag_scale = pag_scale + self._pag_adaptive_scale = pag_adaptive_scale # + # 2. Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1 @@ -810,7 +818,14 @@ def __call__( max_sequence_length=max_sequence_length, ) - if self.do_classifier_free_guidance: + if self.do_perturbed_attention_guidance: + prompt_embeds = self._prepare_perturbed_attention_guidance( + prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance + ) + pooled_prompt_embeds = self._prepare_perturbed_attention_guidance( + pooled_prompt_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance + ) + elif self.do_classifier_free_guidance: prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) @@ -832,14 +847,21 @@ def __call__( latents, ) + if self.do_perturbed_attention_guidance: + original_attn_proc = self.transformer.attn_processors + self._set_pag_attn_processor_sd3( + pag_applied_layers=self.pag_applied_layers, + do_classifier_free_guidance=self.do_classifier_free_guidance, + ) + # 6. Denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: continue - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + # expand the latents if we are doing classifier free guidance, perturbed-attention guidance, or both + latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latent_model_input.shape[0]) @@ -853,7 +875,12 @@ def __call__( )[0] # perform guidance - if self.do_classifier_free_guidance: + if self.do_perturbed_attention_guidance: + noise_pred = self._apply_perturbed_attention_guidance( + noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t + ) + + elif self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) @@ -898,6 +925,9 @@ def __call__( # Offload all models self.maybe_free_model_hooks() + if self.do_perturbed_attention_guidance: + self.transformer.set_attn_processor(original_attn_proc) + if not return_dict: return (image,) From 34a7814564bf00cc07de33c9532d286891755771 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ahn=20Donghoon=20=28=EC=95=88=EB=8F=99=ED=9B=88=20/=20suno?= =?UTF-8?q?=29?= Date: Wed, 17 Jul 2024 14:40:10 +0900 Subject: [PATCH 03/27] Update src/diffusers/models/attention_processor.py Co-authored-by: Sayak Paul --- src/diffusers/models/attention_processor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 8b1d5b8a04ca..bb53cb263ee7 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1079,7 +1079,7 @@ class PAGJointAttnProcessor2_0: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError("PAGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") def __call__( self, From ac1ff12b0270cebf742f0da54454be68bc76c443 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ahn=20Donghoon=20=28=EC=95=88=EB=8F=99=ED=9B=88=20/=20suno?= =?UTF-8?q?=29?= Date: Wed, 17 Jul 2024 14:50:53 +0900 Subject: [PATCH 04/27] Update src/diffusers/models/attention_processor.py Co-authored-by: Sayak Paul --- src/diffusers/models/attention_processor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index bb53cb263ee7..2912aaeedb29 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1224,7 +1224,7 @@ class PAGCFGJointAttnProcessor2_0: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError("PAGCFGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") def __call__( self, From 97b66972e199268585cf0b2a9662747834e3df08 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Sun, 21 Jul 2024 13:19:23 +0900 Subject: [PATCH 05/27] Fix org, ptb, seperate SD3PAGMixion from PAGMixin --- src/diffusers/models/attention_processor.py | 112 +++++++++--------- src/diffusers/pipelines/pag/pag_utils.py | 94 +++++++++++++++ .../pipelines/pag/pipeline_pag_sd_3.py | 8 +- 3 files changed, 154 insertions(+), 60 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 8b1d5b8a04ca..6e5f80d6e596 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1110,9 +1110,9 @@ def __call__( batch_size = encoder_hidden_states_org.shape[0] # `sample` projections. - query = attn.to_q(hidden_states_org) - key = attn.to_k(hidden_states_org) - value = attn.to_v(hidden_states_org) + query_org = attn.to_q(hidden_states_org) + key_org = attn.to_k(hidden_states_org) + value_org = attn.to_v(hidden_states_org) # `context` projections. encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org) @@ -1120,19 +1120,19 @@ def __call__( encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org) # attention - query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) - key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) - value = torch.cat([value, encoder_hidden_states_org_value_proj], dim=1) + query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1) + key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1) + value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1) - inner_dim = key.shape[-1] + inner_dim = key_org.shape[-1] head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - hidden_states_org = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + hidden_states_org = F.scaled_dot_product_attention(query_org, key_org, value_org, dropout_p=0.0, is_causal=False) hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query.dtype) + hidden_states_org = hidden_states_org.to(query_org.dtype) # Split the attention outputs. hidden_states_org, encoder_hidden_states_org = ( @@ -1157,29 +1157,29 @@ def __call__( batch_size = encoder_hidden_states_ptb.shape[0] # `sample` projections. - query = attn.to_q(hidden_states_ptb) - key = attn.to_k(hidden_states_ptb) - value = attn.to_v(hidden_states_ptb) + query_ptb = attn.to_q(hidden_states_ptb) + key_ptb = attn.to_k(hidden_states_ptb) + value_ptb = attn.to_v(hidden_states_ptb) # `context` projections. - encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) - encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) + encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) + encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb) # attention - query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) - key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) - value = torch.cat([value, encoder_hidden_states_ptb_value_proj], dim=1) + query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1) + key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1) + value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1) - inner_dim = key.shape[-1] + inner_dim = key_ptb.shape[-1] head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # create a full mask with all entries set to 0 - seq_len = query.size(2) - full_mask = torch.zeros((seq_len, seq_len), device=query.device, dtype=query.dtype) + seq_len = query_ptb.size(2) + full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype) # set the attention value between image patches to -inf full_mask[:identity_block_size, :identity_block_size] = float('-inf') @@ -1190,9 +1190,9 @@ def __call__( # expand the mask to match the attention weights shape full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions - hidden_states_ptb = F.scaled_dot_product_attention(query, key, value, attn_mask=full_mask, dropout_p=0.0, is_causal=False) + hidden_states_ptb = F.scaled_dot_product_attention(query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False) hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_ptb = hidden_states_ptb.to(query.dtype) + hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype) # split the attention outputs. hidden_states_ptb, encoder_hidden_states_ptb = ( @@ -1259,9 +1259,9 @@ def __call__( batch_size = encoder_hidden_states_org.shape[0] # `sample` projections. - query = attn.to_q(hidden_states_org) - key = attn.to_k(hidden_states_org) - value = attn.to_v(hidden_states_org) + query_org = attn.to_q(hidden_states_org) + key_org = attn.to_k(hidden_states_org) + value_org = attn.to_v(hidden_states_org) # `context` projections. encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org) @@ -1269,19 +1269,19 @@ def __call__( encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org) # attention - query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) - key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) - value = torch.cat([value, encoder_hidden_states_org_value_proj], dim=1) + query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1) + key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1) + value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1) - inner_dim = key.shape[-1] + inner_dim = key_org.shape[-1] head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - hidden_states_org = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + hidden_states_org = F.scaled_dot_product_attention(query_org, key_org, value_org, dropout_p=0.0, is_causal=False) hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query.dtype) + hidden_states_org = hidden_states_org.to(query_org.dtype) # Split the attention outputs. hidden_states_org, encoder_hidden_states_org = ( @@ -1306,29 +1306,29 @@ def __call__( batch_size = encoder_hidden_states_ptb.shape[0] # `sample` projections. - query = attn.to_q(hidden_states_ptb) - key = attn.to_k(hidden_states_ptb) - value = attn.to_v(hidden_states_ptb) + query_ptb = attn.to_q(hidden_states_ptb) + key_ptb = attn.to_k(hidden_states_ptb) + value_ptb = attn.to_v(hidden_states_ptb) # `context` projections. - encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) - encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) + encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) + encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb) # attention - query = torch.cat([query, encoder_hidden_states_org_query_proj], dim=1) - key = torch.cat([key, encoder_hidden_states_org_key_proj], dim=1) - value = torch.cat([value, encoder_hidden_states_ptb_value_proj], dim=1) + query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1) + key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1) + value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1) - inner_dim = key.shape[-1] + inner_dim = key_ptb.shape[-1] head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # create a full mask with all entries set to 0 - seq_len = query.size(2) - full_mask = torch.zeros((seq_len, seq_len), device=query.device, dtype=query.dtype) + seq_len = query_ptb.size(2) + full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype) # set the attention value between image patches to -inf full_mask[:identity_block_size, :identity_block_size] = float('-inf') @@ -1339,9 +1339,9 @@ def __call__( # expand the mask to match the attention weights shape full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions - hidden_states_ptb = F.scaled_dot_product_attention(query, key, value, attn_mask=full_mask, dropout_p=0.0, is_causal=False) + hidden_states_ptb = F.scaled_dot_product_attention(query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False) hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_ptb = hidden_states_ptb.to(query.dtype) + hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype) # split the attention outputs. hidden_states_ptb, encoder_hidden_states_ptb = ( diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py index bf5e2fb87e0e..28b70ff3de10 100644 --- a/src/diffusers/pipelines/pag/pag_utils.py +++ b/src/diffusers/pipelines/pag/pag_utils.py @@ -336,3 +336,97 @@ def pag_attn_processors(self): if proc.__class__ in (PAGCFGIdentitySelfAttnProcessor2_0, PAGIdentitySelfAttnProcessor2_0): processors[name] = proc return processors + +class SD3PAGMixin(PAGMixin): + r"""Mixin class for PAG.""" + + @staticmethod + def _check_input_pag_applied_layer(layer): + r""" + Check if each layer input in `applied_pag_layers` is valid. It should be the block index: {block_index}. + """ + + # Check if the layer index is valid (should be int or str of int) + if isinstance(layer, int): + return # Valid layer index + + if isinstance(layer, str): + if layer.isdigit(): + return # Valid layer index + + # If it is not a valid layer index, raise a ValueError + raise ValueError(f"Pag layer should only contain block index. Accept number string like '3', got {layer}") + + + def _set_pag_attn_processor(self, pag_applied_layers, do_classifier_free_guidance): + r""" + Set the attention processor for the PAG layers. + """ + if do_classifier_free_guidance: + pag_attn_proc = PAGCFGJointAttnProcessor2_0() + else: + pag_attn_proc = PAGJointAttnProcessor2_0() + + def is_attn(module_name): + r""" + Check if the module is self-attention module based on its name. + """ + return "attn" in module_name and len(module_name.split(".")) == 3 # include transformer_blocks.1.attn, exclude transformer_blocks.18.attn.to_q, transformer_blocks.1.attn.add_q_proj, ... + + def get_block_index(module_name): + r""" + Get the block index from the module name. can be "block_0", "block_1", ... If there is only one block (e.g. + mid_block) and index is ommited from the name, it will be "block_0". + """ + # transformer_blocks.23.attn -> "23" + return module_name.split('.')[1] + + for pag_layer_input in pag_applied_layers: + # for each PAG layer input, we find corresponding self-attention layers in the transformer model + target_modules = [] + + block_index = str(pag_layer_input) + + for name, module in self.transformer.named_modules(): + if ( + is_attn(name) + and get_block_index(name) == block_index + ): + target_modules.append(module) + + + if len(target_modules) == 0: + raise ValueError(f"Cannot find pag layer to set attention processor for: {pag_layer_input}") + + for module in target_modules: + module.processor = pag_attn_proc + + + def set_pag_applied_layers(self, pag_applied_layers): + r""" + set the the self-attention layers to apply PAG. Raise ValueError if the input is invalid. + """ + + if not isinstance(pag_applied_layers, list): + pag_applied_layers = [pag_applied_layers] + + for pag_layer in pag_applied_layers: + self._check_input_pag_applied_layer(pag_layer) + + self.pag_applied_layers = pag_applied_layers + + + # TODO: need test + @property + def pag_attn_processors(self): + r""" + Returns: + `dict` of PAG attention processors: A dictionary contains all PAG attention processors used in the model + with the key as the name of the layer. + """ + + processors = {} + for name, proc in self.transformers.attn_processors.items(): + if proc.__class__ in (PAGCFGJointAttnProcessor2_0, PAGJointAttnProcessor2_0): + processors[name] = proc + return processors \ No newline at end of file diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index bb72fc1f1de8..403f84f1ed5b 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -36,7 +36,7 @@ from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput -from .pag_utils import PAGMixin +from .pag_utils import SD3PAGMixin if is_torch_xla_available(): @@ -126,7 +126,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, PAGMixin): +class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3PAGMixin): r""" Args: transformer ([`SD3Transformer2DModel`]): @@ -202,7 +202,7 @@ def __init__( else 128 ) - self.set_pag_applied_layers_sd3(pag_applied_layers) + self.set_pag_applied_layers(pag_applied_layers) def _get_t5_prompt_embeds( @@ -849,7 +849,7 @@ def __call__( if self.do_perturbed_attention_guidance: original_attn_proc = self.transformer.attn_processors - self._set_pag_attn_processor_sd3( + self._set_pag_attn_processor( pag_applied_layers=self.pag_applied_layers, do_classifier_free_guidance=self.do_classifier_free_guidance, ) From efb3ddf6764da5be8fc5da827f221d95105b48c0 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Sun, 21 Jul 2024 13:33:06 +0900 Subject: [PATCH 06/27] Update EXAMPLE_DOC_STRING --- src/diffusers/pipelines/pag/pipeline_pag_sd_3.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index 403f84f1ed5b..0dd8865f9c1a 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -53,15 +53,17 @@ Examples: ```py >>> import torch - >>> from diffusers import StableDiffusion3Pipeline + >>> from diffusers import AutoPipelineForText2Image - >>> pipe = StableDiffusion3Pipeline.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16 + >>> pipe = AutoPipelineForText2Image.from_pretrained( + ... "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, + ... enable_pag=True, + ... pag_applied_layers=["13"] ... ) >>> pipe.to("cuda") >>> prompt = "A cat holding a sign that says hello world" - >>> image = pipe(prompt).images[0] - >>> image.save("sd3.png") + >>> image = pipe(prompt, guidance_scale=5.0, pag_scale=0.7).images[0] + >>> image.save("sd3_pag.png") ``` """ From 67fd57f67ac02d988790b8fb02ac78c6325fcb1c Mon Sep 17 00:00:00 2001 From: sunovivid Date: Sun, 21 Jul 2024 21:16:02 +0900 Subject: [PATCH 07/27] fix missed comma --- src/diffusers/pipelines/auto_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index 81cb0f1d6c9b..5114e1a1fb0d 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -98,7 +98,7 @@ ("stable-diffusion-controlnet-pag", StableDiffusionControlNetPAGPipeline), ("stable-diffusion-xl-pag", StableDiffusionXLPAGPipeline), ("stable-diffusion-xl-controlnet-pag", StableDiffusionXLControlNetPAGPipeline), - ("stable-diffusion-3-medium-diffusers", StableDiffusion3PAGPipeline) + ("stable-diffusion-3-medium-diffusers", StableDiffusion3PAGPipeline), ("auraflow", AuraFlowPipeline), ("kolors", KolorsPipeline), ] From ab8f64283019c467ab57286d4f71e6801b7d98af Mon Sep 17 00:00:00 2001 From: sunovivid Date: Sun, 21 Jul 2024 21:16:20 +0900 Subject: [PATCH 08/27] Remove unused sd3 method in PAGMixin --- src/diffusers/pipelines/pag/pag_utils.py | 76 +----------------------- 1 file changed, 1 insertion(+), 75 deletions(-) diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py index 28b70ff3de10..c86f663b297a 100644 --- a/src/diffusers/pipelines/pag/pag_utils.py +++ b/src/diffusers/pipelines/pag/pag_utils.py @@ -57,24 +57,7 @@ def _check_input_pag_applied_layer(layer): if not layer_splits[2].startswith("attentions_"): raise ValueError(f"Invalid attention_index in pag layer: {layer}. Should start with 'attentions_'") - @staticmethod - def _check_input_pag_applied_layer_sd3(layer): - r""" - Check if each layer input in `applied_pag_layers` is valid. It should be the block index: {block_index}. - """ - - # Check if the layer index is valid (should be int or str of int) - if isinstance(layer, int): - return # Valid layer index - - if isinstance(layer, str): - if layer.isdigit(): - return # Valid layer index - - # If it is not a valid layer index, raise a ValueError - raise ValueError(f"Pag layer should only contain block index. Accept number string like '3', got {layer}") - - + def _set_pag_attn_processor(self, pag_applied_layers, do_classifier_free_guidance): r""" Set the attention processor for the PAG layers. @@ -167,50 +150,6 @@ def get_attn_index(module_name): module.processor = pag_attn_proc - def _set_pag_attn_processor_sd3(self, pag_applied_layers, do_classifier_free_guidance): - r""" - Set the attention processor for the PAG layers. - """ - if do_classifier_free_guidance: - pag_attn_proc = PAGCFGJointAttnProcessor2_0() - else: - pag_attn_proc = PAGJointAttnProcessor2_0() - - def is_attn(module_name): - r""" - Check if the module is self-attention module based on its name. - """ - return "attn" in module_name and len(module_name.split(".")) == 3 # include transformer_blocks.1.attn, exclude transformer_blocks.18.attn.to_q, transformer_blocks.1.attn.add_q_proj, ... - - def get_block_index(module_name): - r""" - Get the block index from the module name. can be "block_0", "block_1", ... If there is only one block (e.g. - mid_block) and index is ommited from the name, it will be "block_0". - """ - # transformer_blocks.23.attn -> "23" - return module_name.split('.')[1] - - for pag_layer_input in pag_applied_layers: - # for each PAG layer input, we find corresponding self-attention layers in the transformer model - target_modules = [] - - block_index = str(pag_layer_input) - - for name, module in self.transformer.named_modules(): - if ( - is_attn(name) - and get_block_index(name) == block_index - ): - target_modules.append(module) - - - if len(target_modules) == 0: - raise ValueError(f"Cannot find pag layer to set attention processor for: {pag_layer_input}") - - for module in target_modules: - module.processor = pag_attn_proc - - def _get_pag_scale(self, t): r""" Get the scale factor for the perturbed attention guidance at timestep `t`. @@ -281,19 +220,7 @@ def set_pag_applied_layers(self, pag_applied_layers): self._check_input_pag_applied_layer(pag_layer) self.pag_applied_layers = pag_applied_layers - - def set_pag_applied_layers_sd3(self, pag_applied_layers): - r""" - set the the self-attention layers to apply PAG. Raise ValueError if the input is invalid. - """ - - if not isinstance(pag_applied_layers, list): - pag_applied_layers = [pag_applied_layers] - - for pag_layer in pag_applied_layers: - self._check_input_pag_applied_layer_sd3(pag_layer) - self.pag_applied_layers = pag_applied_layers @property def pag_scale(self): @@ -416,7 +343,6 @@ def set_pag_applied_layers(self, pag_applied_layers): self.pag_applied_layers = pag_applied_layers - # TODO: need test @property def pag_attn_processors(self): r""" From 795ddafdd936d8783d3cb60b23cf25795dc5f540 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 01:35:16 +0900 Subject: [PATCH 09/27] remove SD3PAGMixin and adopt refactored PAGMixin --- src/diffusers/pipelines/pag/pag_utils.py | 93 ------------------- .../pipelines/pag/pipeline_pag_sd_3.py | 12 ++- 2 files changed, 9 insertions(+), 96 deletions(-) diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py index 0b3d81206180..9096d7187f3b 100644 --- a/src/diffusers/pipelines/pag/pag_utils.py +++ b/src/diffusers/pipelines/pag/pag_utils.py @@ -239,97 +239,4 @@ def pag_attn_processors(self) -> Dict[str, AttentionProcessor]: if proc.__class__ in valid_attn_processors: processors[name] = proc - return processors - -class SD3PAGMixin(PAGMixin): - r"""Mixin class for PAG.""" - - @staticmethod - def _check_input_pag_applied_layer(layer): - r""" - Check if each layer input in `applied_pag_layers` is valid. It should be the block index: {block_index}. - """ - - # Check if the layer index is valid (should be int or str of int) - if isinstance(layer, int): - return # Valid layer index - - if isinstance(layer, str): - if layer.isdigit(): - return # Valid layer index - - # If it is not a valid layer index, raise a ValueError - raise ValueError(f"Pag layer should only contain block index. Accept number string like '3', got {layer}") - - - def _set_pag_attn_processor(self, pag_applied_layers, do_classifier_free_guidance): - r""" - Set the attention processor for the PAG layers. - """ - if do_classifier_free_guidance: - pag_attn_proc = PAGCFGJointAttnProcessor2_0() - else: - pag_attn_proc = PAGJointAttnProcessor2_0() - - def is_attn(module_name): - r""" - Check if the module is self-attention module based on its name. - """ - return "attn" in module_name and len(module_name.split(".")) == 3 # include transformer_blocks.1.attn, exclude transformer_blocks.18.attn.to_q, transformer_blocks.1.attn.add_q_proj, ... - - def get_block_index(module_name): - r""" - Get the block index from the module name. can be "block_0", "block_1", ... If there is only one block (e.g. - mid_block) and index is ommited from the name, it will be "block_0". - """ - # transformer_blocks.23.attn -> "23" - return module_name.split('.')[1] - - for pag_layer_input in pag_applied_layers: - # for each PAG layer input, we find corresponding self-attention layers in the transformer model - target_modules = [] - - block_index = str(pag_layer_input) - - for name, module in self.transformer.named_modules(): - if ( - is_attn(name) - and get_block_index(name) == block_index - ): - target_modules.append(module) - - - if len(target_modules) == 0: - raise ValueError(f"Cannot find pag layer to set attention processor for: {pag_layer_input}") - - for module in target_modules: - module.processor = pag_attn_proc - - - def set_pag_applied_layers(self, pag_applied_layers): - r""" - set the the self-attention layers to apply PAG. Raise ValueError if the input is invalid. - """ - - if not isinstance(pag_applied_layers, list): - pag_applied_layers = [pag_applied_layers] - - for pag_layer in pag_applied_layers: - self._check_input_pag_applied_layer(pag_layer) - - self.pag_applied_layers = pag_applied_layers - - - @property - def pag_attn_processors(self): - r""" - Returns: - `dict` of PAG attention processors: A dictionary contains all PAG attention processors used in the model - with the key as the name of the layer. - """ - - processors = {} - for name, proc in self.transformers.attn_processors.items(): - if proc.__class__ in (PAGCFGJointAttnProcessor2_0, PAGJointAttnProcessor2_0): - processors[name] = proc return processors \ No newline at end of file diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index 0dd8865f9c1a..d469ddb2d1da 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -26,6 +26,7 @@ from ...image_processor import VaeImageProcessor from ...loaders import FromSingleFileMixin, SD3LoraLoaderMixin from ...models.autoencoders import AutoencoderKL +from ...models.attention_processor import PAGCFGJointAttnProcessor2_0, PAGJointAttnProcessor2_0 from ...models.transformers import SD3Transformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( @@ -36,7 +37,7 @@ from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput -from .pag_utils import SD3PAGMixin +from .pag_utils import PAGMixin if is_torch_xla_available(): @@ -128,8 +129,11 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3PAGMixin): +class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, PAGMixin): r""" + [PAG pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag) for text-to-image generation + using Stable Diffusion 3. + Args: transformer ([`SD3Transformer2DModel`]): Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. @@ -204,7 +208,9 @@ def __init__( else 128 ) - self.set_pag_applied_layers(pag_applied_layers) + self.set_pag_applied_layers( + pag_applied_layers, pag_attn_processors=(PAGCFGJointAttnProcessor2_0(), PAGJointAttnProcessor2_0()) + ) def _get_t5_prompt_embeds( From 73df50f595fd7088f501a8d7150ebd13715d6ef0 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 02:27:42 +0900 Subject: [PATCH 10/27] add StableDiffusion3PAGPipeline doc --- docs/source/en/api/pipelines/pag.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/source/en/api/pipelines/pag.md b/docs/source/en/api/pipelines/pag.md index ac12bdb5578d..3007c27b4b09 100644 --- a/docs/source/en/api/pipelines/pag.md +++ b/docs/source/en/api/pipelines/pag.md @@ -74,6 +74,12 @@ Since RegEx is supported as a way for matching layer identifiers, it is crucial - __call__ +## StableDiffusion3PAGPipeline +[[autodoc]] StableDiffusion3PAGPipeline + - all + - __call__ + + ## PixArtSigmaPAGPipeline [[autodoc]] PixArtSigmaPAGPipeline - all From e62232895d3c6f8a36cca354c04bce80709f75fa Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 02:37:44 +0900 Subject: [PATCH 11/27] add parameter description --- src/diffusers/pipelines/pag/pipeline_pag_sd_3.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index d469ddb2d1da..f48537075653 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -181,7 +181,7 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - pag_applied_layers: Union[str, List[str]] = "11", # ["11"] + pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block ): self.register_modules( @@ -756,6 +756,12 @@ def __call__( will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. + pag_scale (`float`, *optional*, defaults to 3.0): + The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention + guidance will not be used. + pag_adaptive_scale (`float`, *optional*, defaults to 0.0): + The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is + used. Examples: From 35349d928349679783ca8e8222324c167f1bcc73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ahn=20Donghoon=20=28=EC=95=88=EB=8F=99=ED=9B=88=20/=20suno?= =?UTF-8?q?=29?= Date: Tue, 6 Aug 2024 02:45:43 +0900 Subject: [PATCH 12/27] remove *args, **kwargs Co-authored-by: YiYi Xu --- src/diffusers/models/attention_processor.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index c7c2a6887e0b..49e7f521099d 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1114,8 +1114,6 @@ def __call__( attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, - *args, - **kwargs, ) -> torch.FloatTensor: residual = hidden_states From 27f1100b83ed71fe6f0c238268ffc4672a724699 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 03:00:20 +0900 Subject: [PATCH 13/27] add import pag pipeline --- src/diffusers/__init__.py | 2 ++ src/diffusers/pipelines/__init__.py | 2 ++ src/diffusers/pipelines/auto_pipeline.py | 2 +- 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index ce437a5e6b73..039950060b0b 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -309,6 +309,7 @@ "StableDiffusion3ControlNetPipeline", "StableDiffusion3Img2ImgPipeline", "StableDiffusion3InpaintPipeline", + "StableDiffusion3PAGPipeline", "StableDiffusion3Pipeline", "StableDiffusionAdapterPipeline", "StableDiffusionAttendAndExcitePipeline", @@ -734,6 +735,7 @@ StableDiffusion3Img2ImgPipeline, StableDiffusion3InpaintPipeline, StableDiffusion3Pipeline, + StableDiffusion3PAGPipeline, StableDiffusionAdapterPipeline, StableDiffusionAttendAndExcitePipeline, StableDiffusionControlNetImg2ImgPipeline, diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 46e93ea7769b..25a32f8cf877 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -146,6 +146,7 @@ [ "AnimateDiffPAGPipeline", "HunyuanDiTPAGPipeline", + "StableDiffusion3PAGPipeline", "StableDiffusionPAGPipeline", "StableDiffusionControlNetPAGPipeline", "StableDiffusionXLPAGPipeline", @@ -535,6 +536,7 @@ AnimateDiffPAGPipeline, HunyuanDiTPAGPipeline, PixArtSigmaPAGPipeline, + StableDiffusion3Pipeline, StableDiffusionControlNetPAGPipeline, StableDiffusionPAGPipeline, StableDiffusionXLControlNetPAGPipeline, diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index a72756b1820e..65f255cecbcd 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -85,6 +85,7 @@ ("stable-diffusion", StableDiffusionPipeline), ("stable-diffusion-xl", StableDiffusionXLPipeline), ("stable-diffusion-3", StableDiffusion3Pipeline), + ("stable-diffusion-3-pag", StableDiffusion3PAGPipeline), ("if", IFPipeline), ("hunyuan", HunyuanDiTPipeline), ("hunyuan-pag", HunyuanDiTPAGPipeline), @@ -102,7 +103,6 @@ ("stable-diffusion-controlnet-pag", StableDiffusionControlNetPAGPipeline), ("stable-diffusion-xl-pag", StableDiffusionXLPAGPipeline), ("stable-diffusion-xl-controlnet-pag", StableDiffusionXLControlNetPAGPipeline), - ("stable-diffusion-3-medium-diffusers", StableDiffusion3PAGPipeline), ("pixart-sigma-pag", PixArtSigmaPAGPipeline), ("auraflow", AuraFlowPipeline), ("kolors", KolorsPipeline), From 512026fab2d0500887da6a80540249b46e07b3b4 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 03:25:22 +0900 Subject: [PATCH 14/27] make style --- src/diffusers/__init__.py | 2 +- src/diffusers/models/attention_processor.py | 6 +++--- src/diffusers/pipelines/auto_pipeline.py | 2 +- src/diffusers/pipelines/pag/__init__.py | 2 +- src/diffusers/pipelines/pag/pag_utils.py | 6 ++---- src/diffusers/pipelines/pag/pipeline_pag_sd_3.py | 14 +++++++------- 6 files changed, 15 insertions(+), 17 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 039950060b0b..131e2c4083dd 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -734,8 +734,8 @@ StableDiffusion3ControlNetPipeline, StableDiffusion3Img2ImgPipeline, StableDiffusion3InpaintPipeline, - StableDiffusion3Pipeline, StableDiffusion3PAGPipeline, + StableDiffusion3Pipeline, StableDiffusionAdapterPipeline, StableDiffusionAttendAndExcitePipeline, StableDiffusionControlNetImg2ImgPipeline, diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 49e7f521099d..4bd1eaa34652 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1179,7 +1179,7 @@ def __call__( encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) ################## perturbed path ################## - + batch_size = encoder_hidden_states_ptb.shape[0] # `sample` projections. @@ -1277,7 +1277,7 @@ def __call__( # chunk hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3) hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org]) - + encoder_hidden_states_uncond, encoder_hidden_states_org, encoder_hidden_states_ptb =encoder_hidden_states.chunk(3) encoder_hidden_states_org = torch.cat([encoder_hidden_states_uncond, encoder_hidden_states_org]) @@ -1328,7 +1328,7 @@ def __call__( encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) ################## perturbed path ################## - + batch_size = encoder_hidden_states_ptb.shape[0] # `sample` projections. diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index 65f255cecbcd..eef00c117d88 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -52,13 +52,13 @@ from .pag import ( HunyuanDiTPAGPipeline, PixArtSigmaPAGPipeline, + StableDiffusion3PAGPipeline, StableDiffusionControlNetPAGPipeline, StableDiffusionPAGPipeline, StableDiffusionXLControlNetPAGPipeline, StableDiffusionXLPAGImg2ImgPipeline, StableDiffusionXLPAGInpaintPipeline, StableDiffusionXLPAGPipeline, - StableDiffusion3PAGPipeline, ) from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline from .stable_cascade import StableCascadeCombinedPipeline, StableCascadeDecoderPipeline diff --git a/src/diffusers/pipelines/pag/__init__.py b/src/diffusers/pipelines/pag/__init__.py index d0b3c3250324..895d6b9f82f8 100644 --- a/src/diffusers/pipelines/pag/__init__.py +++ b/src/diffusers/pipelines/pag/__init__.py @@ -46,11 +46,11 @@ from .pipeline_pag_hunyuandit import HunyuanDiTPAGPipeline from .pipeline_pag_pixart_sigma import PixArtSigmaPAGPipeline from .pipeline_pag_sd import StableDiffusionPAGPipeline + from .pipeline_pag_sd_3 import StableDiffusion3PAGPipeline from .pipeline_pag_sd_animatediff import AnimateDiffPAGPipeline from .pipeline_pag_sd_xl import StableDiffusionXLPAGPipeline from .pipeline_pag_sd_xl_img2img import StableDiffusionXLPAGImg2ImgPipeline from .pipeline_pag_sd_xl_inpaint import StableDiffusionXLPAGInpaintPipeline - from .pipeline_pag_sd_3 import StableDiffusion3PAGPipeline else: import sys diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py index 9096d7187f3b..52ef734b006d 100644 --- a/src/diffusers/pipelines/pag/pag_utils.py +++ b/src/diffusers/pipelines/pag/pag_utils.py @@ -23,8 +23,6 @@ AttentionProcessor, PAGCFGIdentitySelfAttnProcessor2_0, PAGIdentitySelfAttnProcessor2_0, - PAGJointAttnProcessor2_0, - PAGCFGJointAttnProcessor2_0, ) from ...utils import logging @@ -35,7 +33,7 @@ class PAGMixin: r"""Mixin class for [Pertubed Attention Guidance](https://arxiv.org/abs/2403.17377v1).""" - + def _set_pag_attn_processor(self, pag_applied_layers, do_classifier_free_guidance): r""" Set the attention processor for the PAG layers. @@ -239,4 +237,4 @@ def pag_attn_processors(self) -> Dict[str, AttentionProcessor]: if proc.__class__ in valid_attn_processors: processors[name] = proc - return processors \ No newline at end of file + return processors diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index f48537075653..fda0e3722d44 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -25,8 +25,8 @@ from ...image_processor import VaeImageProcessor from ...loaders import FromSingleFileMixin, SD3LoraLoaderMixin -from ...models.autoencoders import AutoencoderKL from ...models.attention_processor import PAGCFGJointAttnProcessor2_0, PAGJointAttnProcessor2_0 +from ...models.autoencoders import AutoencoderKL from ...models.transformers import SD3Transformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( @@ -133,7 +133,7 @@ class StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSin r""" [PAG pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag) for text-to-image generation using Stable Diffusion 3. - + Args: transformer ([`SD3Transformer2DModel`]): Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. @@ -207,7 +207,7 @@ def __init__( if hasattr(self, "transformer") and self.transformer is not None else 128 ) - + self.set_pag_applied_layers( pag_applied_layers, pag_attn_processors=(PAGCFGJointAttnProcessor2_0(), PAGJointAttnProcessor2_0()) ) @@ -797,8 +797,8 @@ def __call__( self._joint_attention_kwargs = joint_attention_kwargs self._interrupt = False self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale # - + self._pag_adaptive_scale = pag_adaptive_scale # + # 2. Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1 @@ -893,7 +893,7 @@ def __call__( noise_pred = self._apply_perturbed_attention_guidance( noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t ) - + elif self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) @@ -941,7 +941,7 @@ def __call__( if self.do_perturbed_attention_guidance: self.transformer.set_attn_processor(original_attn_proc) - + if not return_dict: return (image,) From c0bcc82b0d94adc782d029bd33d6095e882292e1 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 6 Aug 2024 07:15:09 +0530 Subject: [PATCH 15/27] style --- src/diffusers/models/attention_processor.py | 58 ++++++++++++++----- src/diffusers/pipelines/pag/__init__.py | 2 +- src/diffusers/pipelines/pag/pag_utils.py | 3 - .../pipelines/pag/pipeline_pag_sd_3.py | 11 ++-- 4 files changed, 49 insertions(+), 25 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 384266ce334a..133625c16a57 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1111,7 +1111,9 @@ class PAGJointAttnProcessor2_0: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("PAGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError( + "PAGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." + ) def __call__( self, @@ -1130,7 +1132,9 @@ def __call__( batch_size, channel, height, width = encoder_hidden_states.shape encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - identity_block_size = hidden_states.shape[1] # patch embeddings width * height (correspond to self-attention map width or height) + identity_block_size = hidden_states.shape[ + 1 + ] # patch embeddings width * height (correspond to self-attention map width or height) # chunk hidden_states_org, hidden_states_ptb = hidden_states.chunk(2) @@ -1160,7 +1164,9 @@ def __call__( key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - hidden_states_org = F.scaled_dot_product_attention(query_org, key_org, value_org, dropout_p=0.0, is_causal=False) + hidden_states_org = F.scaled_dot_product_attention( + query_org, key_org, value_org, dropout_p=0.0, is_causal=False + ) hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states_org = hidden_states_org.to(query_org.dtype) @@ -1180,7 +1186,9 @@ def __call__( if input_ndim == 4: hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) if context_input_ndim == 4: - encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) + encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) ################## perturbed path ################## @@ -1212,7 +1220,7 @@ def __call__( full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype) # set the attention value between image patches to -inf - full_mask[:identity_block_size, :identity_block_size] = float('-inf') + full_mask[:identity_block_size, :identity_block_size] = float("-inf") # set the diagonal of the attention value between image patches to 0 full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0) @@ -1220,7 +1228,9 @@ def __call__( # expand the mask to match the attention weights shape full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions - hidden_states_ptb = F.scaled_dot_product_attention(query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False) + hidden_states_ptb = F.scaled_dot_product_attention( + query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False + ) hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype) @@ -1240,7 +1250,9 @@ def __call__( if input_ndim == 4: hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) if context_input_ndim == 4: - encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) + encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) ################ concat ############### hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) @@ -1254,7 +1266,9 @@ class PAGCFGJointAttnProcessor2_0: def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("PAGCFGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError( + "PAGCFGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." + ) def __call__( self, @@ -1276,13 +1290,19 @@ def __call__( batch_size, channel, height, width = encoder_hidden_states.shape encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - identity_block_size = hidden_states.shape[1] # patch embeddings width * height (correspond to self-attention map width or height) + identity_block_size = hidden_states.shape[ + 1 + ] # patch embeddings width * height (correspond to self-attention map width or height) # chunk hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3) hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org]) - encoder_hidden_states_uncond, encoder_hidden_states_org, encoder_hidden_states_ptb =encoder_hidden_states.chunk(3) + ( + encoder_hidden_states_uncond, + encoder_hidden_states_org, + encoder_hidden_states_ptb, + ) = encoder_hidden_states.chunk(3) encoder_hidden_states_org = torch.cat([encoder_hidden_states_uncond, encoder_hidden_states_org]) ################## original path ################## @@ -1309,7 +1329,9 @@ def __call__( key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - hidden_states_org = F.scaled_dot_product_attention(query_org, key_org, value_org, dropout_p=0.0, is_causal=False) + hidden_states_org = F.scaled_dot_product_attention( + query_org, key_org, value_org, dropout_p=0.0, is_causal=False + ) hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states_org = hidden_states_org.to(query_org.dtype) @@ -1329,7 +1351,9 @@ def __call__( if input_ndim == 4: hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) if context_input_ndim == 4: - encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) + encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) ################## perturbed path ################## @@ -1361,7 +1385,7 @@ def __call__( full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype) # set the attention value between image patches to -inf - full_mask[:identity_block_size, :identity_block_size] = float('-inf') + full_mask[:identity_block_size, :identity_block_size] = float("-inf") # set the diagonal of the attention value between image patches to 0 full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0) @@ -1369,7 +1393,9 @@ def __call__( # expand the mask to match the attention weights shape full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions - hidden_states_ptb = F.scaled_dot_product_attention(query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False) + hidden_states_ptb = F.scaled_dot_product_attention( + query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False + ) hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype) @@ -1389,7 +1415,9 @@ def __call__( if input_ndim == 4: hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) if context_input_ndim == 4: - encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) + encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) ################ concat ############### hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) diff --git a/src/diffusers/pipelines/pag/__init__.py b/src/diffusers/pipelines/pag/__init__.py index 895d6b9f82f8..8226810c1e6c 100644 --- a/src/diffusers/pipelines/pag/__init__.py +++ b/src/diffusers/pipelines/pag/__init__.py @@ -27,11 +27,11 @@ _import_structure["pipeline_pag_hunyuandit"] = ["HunyuanDiTPAGPipeline"] _import_structure["pipeline_pag_pixart_sigma"] = ["PixArtSigmaPAGPipeline"] _import_structure["pipeline_pag_sd"] = ["StableDiffusionPAGPipeline"] + _import_structure["pipeline_pag_sd_3"] = ["StableDiffusion3PAGPipeline"] _import_structure["pipeline_pag_sd_animatediff"] = ["AnimateDiffPAGPipeline"] _import_structure["pipeline_pag_sd_xl"] = ["StableDiffusionXLPAGPipeline"] _import_structure["pipeline_pag_sd_xl_img2img"] = ["StableDiffusionXLPAGImg2ImgPipeline"] _import_structure["pipeline_pag_sd_xl_inpaint"] = ["StableDiffusionXLPAGInpaintPipeline"] - _import_structure["pipeline_pag_sd_3"] = ["StableDiffusion3PAGPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py index 52ef734b006d..728f730c9904 100644 --- a/src/diffusers/pipelines/pag/pag_utils.py +++ b/src/diffusers/pipelines/pag/pag_utils.py @@ -33,7 +33,6 @@ class PAGMixin: r"""Mixin class for [Pertubed Attention Guidance](https://arxiv.org/abs/2403.17377v1).""" - def _set_pag_attn_processor(self, pag_applied_layers, do_classifier_free_guidance): r""" Set the attention processor for the PAG layers. @@ -86,7 +85,6 @@ def is_fake_integral_match(layer_id, name): for module in target_modules: module.processor = pag_attn_proc - def _get_pag_scale(self, t): r""" Get the scale factor for the perturbed attention guidance at timestep `t`. @@ -189,7 +187,6 @@ def set_pag_applied_layers( self.pag_applied_layers = pag_applied_layers self._pag_attn_processors = pag_attn_processors - @property def pag_scale(self) -> float: r"""Get the scale factor for the perturbed attention guidance.""" diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index fda0e3722d44..53ed742ccd3a 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -57,9 +57,10 @@ >>> from diffusers import AutoPipelineForText2Image >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, + ... "stabilityai/stable-diffusion-3-medium-diffusers", + ... torch_dtype=torch.float16, ... enable_pag=True, - ... pag_applied_layers=["13"] + ... pag_applied_layers=["13"], ... ) >>> pipe.to("cuda") >>> prompt = "A cat holding a sign that says hello world" @@ -181,9 +182,8 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block + pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block ): - self.register_modules( vae=vae, text_encoder=text_encoder, @@ -212,7 +212,6 @@ def __init__( pag_applied_layers, pag_attn_processors=(PAGCFGJointAttnProcessor2_0(), PAGJointAttnProcessor2_0()) ) - def _get_t5_prompt_embeds( self, prompt: Union[str, List[str]] = None, @@ -797,7 +796,7 @@ def __call__( self._joint_attention_kwargs = joint_attention_kwargs self._interrupt = False self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale # + self._pag_adaptive_scale = pag_adaptive_scale # # 2. Define call parameters if prompt is not None and isinstance(prompt, str): From 12c49dbf19d0133082d7e9b1674514e25bf44122 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 14:27:15 +0900 Subject: [PATCH 16/27] apply comments (copied from) --- src/diffusers/pipelines/pag/pipeline_pag_sd_3.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index 53ed742ccd3a..484388e6bde6 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -60,7 +60,7 @@ ... "stabilityai/stable-diffusion-3-medium-diffusers", ... torch_dtype=torch.float16, ... enable_pag=True, - ... pag_applied_layers=["13"], + ... pag_applied_layers=["blocks.13"], ... ) >>> pipe.to("cuda") >>> prompt = "A cat holding a sign that says hello world" @@ -212,6 +212,7 @@ def __init__( pag_applied_layers, pag_attn_processors=(PAGCFGJointAttnProcessor2_0(), PAGJointAttnProcessor2_0()) ) + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, prompt: Union[str, List[str]] = None, @@ -268,6 +269,7 @@ def _get_t5_prompt_embeds( return prompt_embeds + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, prompt: Union[str, List[str]], @@ -323,6 +325,7 @@ def _get_clip_prompt_embeds( return prompt_embeds, pooled_prompt_embeds + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt def encode_prompt( self, prompt: Union[str, List[str]], @@ -492,6 +495,7 @@ def encode_prompt( return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.check_inputs def check_inputs( self, prompt, @@ -582,6 +586,7 @@ def check_inputs( 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}") + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents def prepare_latents( self, batch_size, From 1c320604300bc25b93f7d367aba015e4316b90cc Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 14:27:39 +0900 Subject: [PATCH 17/27] apply comments (fix the comment location) --- src/diffusers/models/attention_processor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 133625c16a57..e2ab1606b345 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -1132,9 +1132,9 @@ def __call__( batch_size, channel, height, width = encoder_hidden_states.shape encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - identity_block_size = hidden_states.shape[ - 1 - ] # patch embeddings width * height (correspond to self-attention map width or height) + # store the length of image patch sequences to create a mask that prevents interaction between patches + # similar to making the self-attention map an identity matrix + identity_block_size = hidden_states.shape[1] # chunk hidden_states_org, hidden_states_ptb = hidden_states.chunk(2) From fa90a92b96dcc95fe640b08e782390f0d5138f9a Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 14:33:35 +0900 Subject: [PATCH 18/27] make fix-copies, sync SD3 pipeline update --- .../pipelines/pag/pipeline_pag_sd_3.py | 35 ++++++++++++++++++- .../dummy_torch_and_transformers_objects.py | 15 ++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index 484388e6bde6..3035509843c0 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -30,9 +30,12 @@ from ...models.transformers import SD3Transformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( + USE_PEFT_BACKEND, is_torch_xla_available, logging, replace_example_docstring, + scale_lora_layers, + unscale_lora_layers, ) from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline @@ -184,6 +187,8 @@ def __init__( tokenizer_3: T5TokenizerFast, pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block ): + super().__init__() + self.register_modules( vae=vae, text_encoder=text_encoder, @@ -343,6 +348,7 @@ def encode_prompt( negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, clip_skip: Optional[int] = None, max_sequence_length: int = 256, + lora_scale: Optional[float] = None, ): r""" @@ -388,9 +394,22 @@ def encode_prompt( clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that the output of the pre-final layer will be used for computing the prompt embeddings. + lora_scale (`float`, *optional*): + A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. """ device = device or self._execution_device + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, SD3LoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if self.text_encoder is not None and USE_PEFT_BACKEND: + scale_lora_layers(self.text_encoder, lora_scale) + if self.text_encoder_2 is not None and USE_PEFT_BACKEND: + scale_lora_layers(self.text_encoder_2, lora_scale) + prompt = [prompt] if isinstance(prompt, str) else prompt if prompt is not None: batch_size = len(prompt) @@ -493,6 +512,16 @@ def encode_prompt( [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 ) + if self.text_encoder is not None: + if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder, lora_scale) + + if self.text_encoder_2 is not None: + if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder_2, lora_scale) + return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.check_inputs @@ -701,7 +730,7 @@ def __call__( Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - guidance_scale (`float`, *optional*, defaults to 5.0): + guidance_scale (`float`, *optional*, defaults to 7.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > @@ -813,6 +842,9 @@ def __call__( device = self._execution_device + lora_scale = ( + self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ) ( prompt_embeds, negative_prompt_embeds, @@ -834,6 +866,7 @@ def __call__( clip_skip=self.clip_skip, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + lora_scale=lora_scale, ) if self.do_perturbed_attention_guidance: diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index ad3a1663daa2..3512d09cb57d 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1127,6 +1127,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class StableDiffusion3PAGPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class StableDiffusion3Pipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From 986bcefc480c91c515e7674e0e918251b15b9fd2 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 14:43:59 +0900 Subject: [PATCH 19/27] add test --- tests/pipelines/pag/test_pag_sd3.py | 361 ++++++++++++++++++++++++++++ 1 file changed, 361 insertions(+) create mode 100644 tests/pipelines/pag/test_pag_sd3.py diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py new file mode 100644 index 000000000000..c6115c60a914 --- /dev/null +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -0,0 +1,361 @@ +import inspect +import gc +import unittest + +import numpy as np +import torch +from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel + +from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline +from diffusers.utils.testing_utils import ( + numpy_cosine_similarity_distance, + require_torch_gpu, + slow, + torch_device, +) + +from ..test_pipelines_common import ( + PipelineTesterMixin, + check_qkv_fusion_matches_attn_procs_length, + check_qkv_fusion_processors_exist, +) + + +class StableDiffusion3PipelineFastTests(unittest.TestCase, PipelineTesterMixin): + pipeline_class = StableDiffusion3Pipeline + params = frozenset( + [ + "prompt", + "height", + "width", + "guidance_scale", + "negative_prompt", + "prompt_embeds", + "negative_prompt_embeds", + ] + ) + batch_params = frozenset(["prompt", "negative_prompt"]) + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = SD3Transformer2DModel( + sample_size=32, + patch_size=1, + in_channels=4, + num_layers=1, + attention_head_dim=8, + num_attention_heads=4, + caption_projection_dim=32, + joint_attention_dim=32, + pooled_projection_dim=64, + out_channels=4, + ) + clip_text_encoder_config = CLIPTextConfig( + bos_token_id=0, + eos_token_id=2, + hidden_size=32, + intermediate_size=37, + layer_norm_eps=1e-05, + num_attention_heads=4, + num_hidden_layers=5, + pad_token_id=1, + vocab_size=1000, + hidden_act="gelu", + projection_dim=32, + ) + + torch.manual_seed(0) + text_encoder = CLIPTextModelWithProjection(clip_text_encoder_config) + + torch.manual_seed(0) + text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) + + text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + + tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + tokenizer_3 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + torch.manual_seed(0) + vae = AutoencoderKL( + sample_size=32, + in_channels=3, + out_channels=3, + block_out_channels=(4,), + layers_per_block=1, + latent_channels=4, + norm_num_groups=1, + use_quant_conv=False, + use_post_quant_conv=False, + shift_factor=0.0609, + scaling_factor=1.5035, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "scheduler": scheduler, + "text_encoder": text_encoder, + "text_encoder_2": text_encoder_2, + "text_encoder_3": text_encoder_3, + "tokenizer": tokenizer, + "tokenizer_2": tokenizer_2, + "tokenizer_3": tokenizer_3, + "transformer": transformer, + "vae": vae, + } + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device="cpu").manual_seed(seed) + + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "output_type": "np", + } + return inputs + + def test_stable_diffusion_3_different_prompts(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + + inputs = self.get_dummy_inputs(torch_device) + output_same_prompt = pipe(**inputs).images[0] + + inputs = self.get_dummy_inputs(torch_device) + inputs["prompt_2"] = "a different prompt" + inputs["prompt_3"] = "another different prompt" + output_different_prompts = pipe(**inputs).images[0] + + max_diff = np.abs(output_same_prompt - output_different_prompts).max() + + # Outputs should be different here + assert max_diff > 1e-2 + + def test_stable_diffusion_3_different_negative_prompts(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + + inputs = self.get_dummy_inputs(torch_device) + output_same_prompt = pipe(**inputs).images[0] + + inputs = self.get_dummy_inputs(torch_device) + inputs["negative_prompt_2"] = "deformed" + inputs["negative_prompt_3"] = "blurry" + output_different_prompts = pipe(**inputs).images[0] + + max_diff = np.abs(output_same_prompt - output_different_prompts).max() + + # Outputs should be different here + assert max_diff > 1e-2 + + def test_stable_diffusion_3_prompt_embeds(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + + output_with_prompt = pipe(**inputs).images[0] + + inputs = self.get_dummy_inputs(torch_device) + prompt = inputs.pop("prompt") + + do_classifier_free_guidance = inputs["guidance_scale"] > 1 + ( + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = pipe.encode_prompt( + prompt, + prompt_2=None, + prompt_3=None, + do_classifier_free_guidance=do_classifier_free_guidance, + device=torch_device, + ) + output_with_embeds = pipe( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + **inputs, + ).images[0] + + max_diff = np.abs(output_with_prompt - output_with_embeds).max() + assert max_diff < 1e-4 + + def test_fused_qkv_projections(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + original_image_slice = image[0, -3:, -3:, -1] + + # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added + # to the pipeline level. + pipe.transformer.fuse_qkv_projections() + assert check_qkv_fusion_processors_exist( + pipe.transformer + ), "Something wrong with the fused attention processors. Expected all the attention processors to be fused." + assert check_qkv_fusion_matches_attn_procs_length( + pipe.transformer, pipe.transformer.original_attn_processors + ), "Something wrong with the attention processors concerning the fused QKV projections." + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + image_slice_fused = image[0, -3:, -3:, -1] + + pipe.transformer.unfuse_qkv_projections() + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + image_slice_disabled = image[0, -3:, -3:, -1] + + assert np.allclose( + original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3 + ), "Fusion of QKV projections shouldn't affect the outputs." + assert np.allclose( + image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3 + ), "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert np.allclose( + original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2 + ), "Original outputs should match when fused QKV projections are disabled." + + def test_pag_disable_enable(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components() + + # base pipeline (expect same output when pag is disabled) + pipe_sd = StableDiffusion3Pipeline(**components) + pipe_sd = pipe_sd.to(device) + pipe_sd.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + del inputs["pag_scale"] + assert ( + "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters + ), f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." + out = pipe_sd(**inputs).images[0, -3:, -3:, -1] + + components = self.get_dummy_components() + + # pag disabled with pag_scale=0.0 + pipe_pag = self.pipeline_class(**components) + pipe_pag = pipe_pag.to(device) + pipe_pag.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["pag_scale"] = 0.0 + out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] + + # pag enabled + pipe_pag = self.pipeline_class(**components) + pipe_pag = pipe_pag.to(device) + pipe_pag.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["pag_scale"] = 3.0 + out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] + + assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 + assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 + + def test_pag_applied_layers(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components() + + # base pipeline + pipe = self.pipeline_class(**components) + pipe = pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k] + original_attn_procs = pipe.transformer.attn_processors + pag_layers = ["blocks.0", "blocks.1"] + pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) + assert set(pipe.pag_attn_processors) == set(all_self_attn_layers) + + # blocks.0 + block_0_self_attn = ["blocks.0.attn1.processor"] + pipe.transformer.set_attn_processor(original_attn_procs.copy()) + pag_layers = ["blocks.0"] + pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) + assert set(pipe.pag_attn_processors) == set(block_0_self_attn) + + pipe.transformer.set_attn_processor(original_attn_procs.copy()) + pag_layers = ["blocks.0.attn1"] + pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) + assert set(pipe.pag_attn_processors) == set(block_0_self_attn) + + pipe.transformer.set_attn_processor(original_attn_procs.copy()) + pag_layers = ["blocks.(0|1)"] + pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) + assert (len(pipe.pag_attn_processors)) == 2 + + pipe.transformer.set_attn_processor(original_attn_procs.copy()) + pag_layers = ["blocks.0", r"blocks\.1"] + pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) + assert len(pipe.pag_attn_processors) == 2 + + +@slow +@require_torch_gpu +class StableDiffusion3PipelineSlowTests(unittest.TestCase): + pipeline_class = StableDiffusion3Pipeline + repo_id = "stabilityai/stable-diffusion-3-medium-diffusers" + + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + + def tearDown(self): + super().tearDown() + gc.collect() + torch.cuda.empty_cache() + + def get_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device="cpu").manual_seed(seed) + + return { + "prompt": "A photo of a cat", + "num_inference_steps": 2, + "guidance_scale": 5.0, + "output_type": "np", + "generator": generator, + } + + def test_sd3_inference(self): + pipe = self.pipeline_class.from_pretrained(self.repo_id, torch_dtype=torch.float16) + pipe.enable_model_cpu_offload() + + inputs = self.get_inputs(torch_device) + + image = pipe(**inputs).images[0] + image_slice = image[0, :10, :10] + expected_slice = np.array( + [ + [0.36132812, 0.30004883, 0.25830078], + [0.36669922, 0.31103516, 0.23754883], + [0.34814453, 0.29248047, 0.23583984], + [0.35791016, 0.30981445, 0.23999023], + [0.36328125, 0.31274414, 0.2607422], + [0.37304688, 0.32177734, 0.26171875], + [0.3671875, 0.31933594, 0.25756836], + [0.36035156, 0.31103516, 0.2578125], + [0.3857422, 0.33789062, 0.27563477], + [0.3701172, 0.31982422, 0.265625], + ], + dtype=np.float32, + ) + + max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten()) + + assert max_diff < 1e-4 \ No newline at end of file From ee863fa94ebb85ca0ffb839cf67e6d9c6889016a Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 14:49:17 +0900 Subject: [PATCH 20/27] fix test following HunyuanDiT test -> HunyuanDiTPAG test --- tests/pipelines/pag/test_pag_sd3.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index c6115c60a914..06feb6d0c2b8 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -6,7 +6,7 @@ import torch from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel -from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline +from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline, StableDiffusion3PAGPipeline from diffusers.utils.testing_utils import ( numpy_cosine_similarity_distance, require_torch_gpu, @@ -21,8 +21,8 @@ ) -class StableDiffusion3PipelineFastTests(unittest.TestCase, PipelineTesterMixin): - pipeline_class = StableDiffusion3Pipeline +class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixin): + pipeline_class = StableDiffusion3PAGPipeline params = frozenset( [ "prompt", @@ -117,6 +117,7 @@ def get_dummy_inputs(self, device, seed=0): "num_inference_steps": 2, "guidance_scale": 5.0, "output_type": "np", + "pag_scale": 0.0, } return inputs @@ -304,8 +305,8 @@ def test_pag_applied_layers(self): @slow @require_torch_gpu -class StableDiffusion3PipelineSlowTests(unittest.TestCase): - pipeline_class = StableDiffusion3Pipeline +class StableDiffusion3PAGPipelineSlowTests(unittest.TestCase): + pipeline_class = StableDiffusion3PAGPipeline repo_id = "stabilityai/stable-diffusion-3-medium-diffusers" def setUp(self): From e243d5d7f55ef10bdb60dbd7b058cb1037e4e1d2 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 14:50:13 +0900 Subject: [PATCH 21/27] remove slow test --- tests/pipelines/pag/test_pag_sd3.py | 61 +---------------------------- 1 file changed, 1 insertion(+), 60 deletions(-) diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 06feb6d0c2b8..461dd2e5b0b4 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -300,63 +300,4 @@ def test_pag_applied_layers(self): pipe.transformer.set_attn_processor(original_attn_procs.copy()) pag_layers = ["blocks.0", r"blocks\.1"] pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) - assert len(pipe.pag_attn_processors) == 2 - - -@slow -@require_torch_gpu -class StableDiffusion3PAGPipelineSlowTests(unittest.TestCase): - pipeline_class = StableDiffusion3PAGPipeline - repo_id = "stabilityai/stable-diffusion-3-medium-diffusers" - - def setUp(self): - super().setUp() - gc.collect() - torch.cuda.empty_cache() - - def tearDown(self): - super().tearDown() - gc.collect() - torch.cuda.empty_cache() - - def get_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - - return { - "prompt": "A photo of a cat", - "num_inference_steps": 2, - "guidance_scale": 5.0, - "output_type": "np", - "generator": generator, - } - - def test_sd3_inference(self): - pipe = self.pipeline_class.from_pretrained(self.repo_id, torch_dtype=torch.float16) - pipe.enable_model_cpu_offload() - - inputs = self.get_inputs(torch_device) - - image = pipe(**inputs).images[0] - image_slice = image[0, :10, :10] - expected_slice = np.array( - [ - [0.36132812, 0.30004883, 0.25830078], - [0.36669922, 0.31103516, 0.23754883], - [0.34814453, 0.29248047, 0.23583984], - [0.35791016, 0.30981445, 0.23999023], - [0.36328125, 0.31274414, 0.2607422], - [0.37304688, 0.32177734, 0.26171875], - [0.3671875, 0.31933594, 0.25756836], - [0.36035156, 0.31103516, 0.2578125], - [0.3857422, 0.33789062, 0.27563477], - [0.3701172, 0.31982422, 0.265625], - ], - dtype=np.float32, - ) - - max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten()) - - assert max_diff < 1e-4 \ No newline at end of file + assert len(pipe.pag_attn_processors) == 2 \ No newline at end of file From ca33100785e47e8198450546e6b62af22cad5fb5 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 14:51:32 +0900 Subject: [PATCH 22/27] make style --- tests/pipelines/pag/test_pag_sd3.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 461dd2e5b0b4..0a8869e0a8d9 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -1,16 +1,18 @@ import inspect -import gc import unittest import numpy as np import torch from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel -from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline, StableDiffusion3PAGPipeline +from diffusers import ( + AutoencoderKL, + FlowMatchEulerDiscreteScheduler, + SD3Transformer2DModel, + StableDiffusion3PAGPipeline, + StableDiffusion3Pipeline, +) from diffusers.utils.testing_utils import ( - numpy_cosine_similarity_distance, - require_torch_gpu, - slow, torch_device, ) @@ -300,4 +302,4 @@ def test_pag_applied_layers(self): pipe.transformer.set_attn_processor(original_attn_procs.copy()) pag_layers = ["blocks.0", r"blocks\.1"] pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) - assert len(pipe.pag_attn_processors) == 2 \ No newline at end of file + assert len(pipe.pag_attn_processors) == 2 From 12ec17c751269b584a29fa6f794014595ab392f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ahn=20Donghoon=20=28=EC=95=88=EB=8F=99=ED=9B=88=20/=20suno?= =?UTF-8?q?=29?= Date: Tue, 6 Aug 2024 15:25:33 +0900 Subject: [PATCH 23/27] Update src/diffusers/pipelines/__init__.py Co-authored-by: Aryan --- src/diffusers/pipelines/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 7d9cc12e21d7..2b34d1935e79 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -541,7 +541,7 @@ AnimateDiffPAGPipeline, HunyuanDiTPAGPipeline, PixArtSigmaPAGPipeline, - StableDiffusion3Pipeline, + StableDiffusion3PAGPipeline, StableDiffusionControlNetPAGPipeline, StableDiffusionPAGPipeline, StableDiffusionXLControlNetPAGPipeline, From 78867bb84215759008cb214a8427e8b8172522e7 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 15:39:36 +0900 Subject: [PATCH 24/27] set num_layers=2 for sd3 models in SD3 PAG test --- tests/pipelines/pag/test_pag_sd3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 0a8869e0a8d9..702989438634 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -44,7 +44,7 @@ def get_dummy_components(self): sample_size=32, patch_size=1, in_channels=4, - num_layers=1, + num_layers=2, attention_head_dim=8, num_attention_heads=4, caption_projection_dim=32, From c6affa7bd2315e505c87674a36043872a09003d7 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 17:38:10 +0900 Subject: [PATCH 25/27] fix attn1 to attn --- tests/pipelines/pag/test_pag_sd3.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 702989438634..e6b3f61f37b4 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -276,21 +276,21 @@ def test_pag_applied_layers(self): pipe = pipe.to(device) pipe.set_progress_bar_config(disable=None) - all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k] + all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn" in k] original_attn_procs = pipe.transformer.attn_processors pag_layers = ["blocks.0", "blocks.1"] pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert set(pipe.pag_attn_processors) == set(all_self_attn_layers) # blocks.0 - block_0_self_attn = ["blocks.0.attn1.processor"] + block_0_self_attn = ["transformer_blocks.0.attn.processor"] pipe.transformer.set_attn_processor(original_attn_procs.copy()) pag_layers = ["blocks.0"] pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert set(pipe.pag_attn_processors) == set(block_0_self_attn) pipe.transformer.set_attn_processor(original_attn_procs.copy()) - pag_layers = ["blocks.0.attn1"] + pag_layers = ["blocks.0.attn"] pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert set(pipe.pag_attn_processors) == set(block_0_self_attn) From 93b03497ec7209b2fef524327ddd0e12ee827a56 Mon Sep 17 00:00:00 2001 From: sunovivid Date: Tue, 6 Aug 2024 23:37:04 +0900 Subject: [PATCH 26/27] remove assert case for nan --- tests/pipelines/pag/test_pag_sd3.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index e6b3f61f37b4..8fd5c6542d82 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -265,7 +265,6 @@ def test_pag_disable_enable(self): out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 def test_pag_applied_layers(self): device = "cpu" # ensure determinism for the device-dependent torch.Generator From 792398de5491b29cd2bb8c37b67302f59b8012f9 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 6 Aug 2024 20:52:32 +0200 Subject: [PATCH 27/27] style --- tests/pipelines/pag/test_pag_sd3.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 8fd5c6542d82..93260870e723 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -255,15 +255,6 @@ def test_pag_disable_enable(self): inputs["pag_scale"] = 0.0 out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - # pag enabled - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 3.0 - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 def test_pag_applied_layers(self):