diffusers/community-pipelines-mirror

1# Based on stable_diffusion_xl_reference.py and stable_diffusion_controlnet_reference.py2 3import inspect4from typing import Any, Callable, Dict, List, Optional, Tuple, Union5 6import numpy as np7import PIL.Image8import torch9 10from diffusers import StableDiffusionXLControlNetPipeline11from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback12from diffusers.image_processor import PipelineImageInput13from diffusers.models import ControlNetModel14from diffusers.models.attention import BasicTransformerBlock15from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, UpBlock2D16from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel17from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput18from diffusers.utils import PIL_INTERPOLATION, deprecate, logging, replace_example_docstring19from diffusers.utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor20 21 22logger = logging.get_logger(__name__)  # pylint: disable=invalid-name23 24 25EXAMPLE_DOC_STRING = """26    Examples:27        ```py28        >>> # !pip install opencv-python transformers accelerate29        >>> from diffusers import ControlNetModel, AutoencoderKL30        >>> from diffusers.schedulers import UniPCMultistepScheduler31        >>> from diffusers.utils import load_image32        >>> import numpy as np33        >>> import torch34 35        >>> import cv236        >>> from PIL import Image37 38        >>> # download an image for the Canny controlnet39        >>> canny_image = load_image(40        ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl_reference_input_cat.jpg"41        ... )42 43        >>> # download an image for the Reference controlnet44        >>> ref_image = load_image(45        ...     "https://hf.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png"46        ... )47 48        >>> # initialize the models and pipeline49        >>> controlnet_conditioning_scale = 0.5  # recommended for good generalization50        >>> controlnet = ControlNetModel.from_pretrained(51        ...     "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float1652        ... )53        >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)54        >>> pipe = StableDiffusionXLControlNetReferencePipeline.from_pretrained(55        ...     "stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnet, vae=vae, torch_dtype=torch.float1656        ... ).to("cuda:0")57 58        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)59 60        >>> # get canny image61        >>> image = np.array(canny_image)62        >>> image = cv2.Canny(image, 100, 200)63        >>> image = image[:, :, None]64        >>> image = np.concatenate([image, image, image], axis=2)65        >>> canny_image = Image.fromarray(image)66 67        >>> # generate image68        >>> image = pipe(69        ...     prompt="a cat",70        ...     num_inference_steps=20,71        ...     controlnet_conditioning_scale=controlnet_conditioning_scale,72        ...     image=canny_image,73        ...     ref_image=ref_image,74        ...     reference_attn=True,75        ...     reference_adain=True76        ...     style_fidelity=1.0,77        ...     generator=torch.Generator("cuda").manual_seed(42)78        ... ).images[0]79        ```80"""81 82 83def torch_dfs(model: torch.nn.Module):84    result = [model]85    for child in model.children():86        result += torch_dfs(child)87    return result88 89 90# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps91def retrieve_timesteps(92    scheduler,93    num_inference_steps: Optional[int] = None,94    device: Optional[Union[str, torch.device]] = None,95    timesteps: Optional[List[int]] = None,96    sigmas: Optional[List[float]] = None,97    **kwargs,98):99    r"""100    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles101    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.102 103    Args:104        scheduler (`SchedulerMixin`):105            The scheduler to get timesteps from.106        num_inference_steps (`int`):107            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`108            must be `None`.109        device (`str` or `torch.device`, *optional*):110            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.111        timesteps (`List[int]`, *optional*):112            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,113            `num_inference_steps` and `sigmas` must be `None`.114        sigmas (`List[float]`, *optional*):115            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,116            `num_inference_steps` and `timesteps` must be `None`.117 118    Returns:119        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the120        second element is the number of inference steps.121    """122    if timesteps is not None and sigmas is not None:123        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")124    if timesteps is not None:125        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())126        if not accepts_timesteps:127            raise ValueError(128                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"129                f" timestep schedules. Please check whether you are using the correct scheduler."130            )131        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)132        timesteps = scheduler.timesteps133        num_inference_steps = len(timesteps)134    elif sigmas is not None:135        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())136        if not accept_sigmas:137            raise ValueError(138                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"139                f" sigmas schedules. Please check whether you are using the correct scheduler."140            )141        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)142        timesteps = scheduler.timesteps143        num_inference_steps = len(timesteps)144    else:145        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)146        timesteps = scheduler.timesteps147    return timesteps, num_inference_steps148 149 150class StableDiffusionXLControlNetReferencePipeline(StableDiffusionXLControlNetPipeline):151    r"""152    Pipeline for text-to-image generation using Stable Diffusion XL with ControlNet guidance.153 154    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods155    implemented for all pipelines (downloading, saving, running on a particular device, etc.).156 157    The pipeline also inherits the following loading methods:158        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings159        - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights160        - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights161        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files162        - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters163 164    Args:165        vae ([`AutoencoderKL`]):166            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.167        text_encoder ([`~transformers.CLIPTextModel`]):168            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).169        text_encoder_2 ([`~transformers.CLIPTextModelWithProjection`]):170            Second frozen text-encoder171            ([laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)).172        tokenizer ([`~transformers.CLIPTokenizer`]):173            A `CLIPTokenizer` to tokenize text.174        tokenizer_2 ([`~transformers.CLIPTokenizer`]):175            A `CLIPTokenizer` to tokenize text.176        unet ([`UNet2DConditionModel`]):177            A `UNet2DConditionModel` to denoise the encoded image latents.178        controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):179            Provides additional conditioning to the `unet` during the denoising process. If you set multiple180            ControlNets as a list, the outputs from each ControlNet are added together to create one combined181            additional conditioning.182        scheduler ([`SchedulerMixin`]):183            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of184            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].185        force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):186            Whether the negative prompt embeddings should always be set to 0. Also see the config of187            `stabilityai/stable-diffusion-xl-base-1-0`.188        add_watermarker (`bool`, *optional*):189            Whether to use the [invisible_watermark](https://github.com/ShieldMnt/invisible-watermark/) library to190            watermark output images. If not defined, it defaults to `True` if the package is installed; otherwise no191            watermarker is used.192    """193 194    def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):195        refimage = refimage.to(device=device)196        if self.vae.dtype == torch.float16 and self.vae.config.force_upcast:197            self.upcast_vae()198            refimage = refimage.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)199        if refimage.dtype != self.vae.dtype:200            refimage = refimage.to(dtype=self.vae.dtype)201        # encode the mask image into latents space so we can concatenate it to the latents202        if isinstance(generator, list):203            ref_image_latents = [204                self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])205                for i in range(batch_size)206            ]207            ref_image_latents = torch.cat(ref_image_latents, dim=0)208        else:209            ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)210        ref_image_latents = self.vae.config.scaling_factor * ref_image_latents211 212        # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method213        if ref_image_latents.shape[0] < batch_size:214            if not batch_size % ref_image_latents.shape[0] == 0:215                raise ValueError(216                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"217                    f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."218                    " Make sure the number of images that you pass is divisible by the total requested batch size."219                )220            ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)221 222        ref_image_latents = torch.cat([ref_image_latents] * 2) if do_classifier_free_guidance else ref_image_latents223 224        # aligning device to prevent device errors when concating it with the latent model input225        ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)226        return ref_image_latents227 228    def prepare_ref_image(229        self,230        image,231        width,232        height,233        batch_size,234        num_images_per_prompt,235        device,236        dtype,237        do_classifier_free_guidance=False,238        guess_mode=False,239    ):240        if not isinstance(image, torch.Tensor):241            if isinstance(image, PIL.Image.Image):242                image = [image]243 244            if isinstance(image[0], PIL.Image.Image):245                images = []246 247                for image_ in image:248                    image_ = image_.convert("RGB")249                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])250                    image_ = np.array(image_)251                    image_ = image_[None, :]252                    images.append(image_)253 254                image = images255 256                image = np.concatenate(image, axis=0)257                image = np.array(image).astype(np.float32) / 255.0258                image = (image - 0.5) / 0.5259                image = image.transpose(0, 3, 1, 2)260                image = torch.from_numpy(image)261 262            elif isinstance(image[0], torch.Tensor):263                image = torch.stack(image, dim=0)264 265        image_batch_size = image.shape[0]266 267        if image_batch_size == 1:268            repeat_by = batch_size269        else:270            repeat_by = num_images_per_prompt271 272        image = image.repeat_interleave(repeat_by, dim=0)273 274        image = image.to(device=device, dtype=dtype)275 276        if do_classifier_free_guidance and not guess_mode:277            image = torch.cat([image] * 2)278 279        return image280 281    def check_ref_inputs(282        self,283        ref_image,284        reference_guidance_start,285        reference_guidance_end,286        style_fidelity,287        reference_attn,288        reference_adain,289    ):290        ref_image_is_pil = isinstance(ref_image, PIL.Image.Image)291        ref_image_is_tensor = isinstance(ref_image, torch.Tensor)292 293        if not ref_image_is_pil and not ref_image_is_tensor:294            raise TypeError(295                f"ref image must be passed and be one of PIL image or torch tensor, but is {type(ref_image)}"296            )297 298        if not reference_attn and not reference_adain:299            raise ValueError("`reference_attn` or `reference_adain` must be True.")300 301        if style_fidelity < 0.0:302            raise ValueError(f"style fidelity: {style_fidelity} can't be smaller than 0.")303        if style_fidelity > 1.0:304            raise ValueError(f"style fidelity: {style_fidelity} can't be larger than 1.0.")305 306        if reference_guidance_start >= reference_guidance_end:307            raise ValueError(308                f"reference guidance start: {reference_guidance_start} cannot be larger or equal to reference guidance end: {reference_guidance_end}."309            )310        if reference_guidance_start < 0.0:311            raise ValueError(f"reference guidance start: {reference_guidance_start} can't be smaller than 0.")312        if reference_guidance_end > 1.0:313            raise ValueError(f"reference guidance end: {reference_guidance_end} can't be larger than 1.0.")314 315    @torch.no_grad()316    @replace_example_docstring(EXAMPLE_DOC_STRING)317    def __call__(318        self,319        prompt: Union[str, List[str]] = None,320        prompt_2: Optional[Union[str, List[str]]] = None,321        image: PipelineImageInput = None,322        ref_image: Union[torch.Tensor, PIL.Image.Image] = None,323        height: Optional[int] = None,324        width: Optional[int] = None,325        num_inference_steps: int = 50,326        timesteps: List[int] = None,327        sigmas: List[float] = None,328        denoising_end: Optional[float] = None,329        guidance_scale: float = 5.0,330        negative_prompt: Optional[Union[str, List[str]]] = None,331        negative_prompt_2: Optional[Union[str, List[str]]] = None,332        num_images_per_prompt: Optional[int] = 1,333        eta: float = 0.0,334        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,335        latents: Optional[torch.Tensor] = None,336        prompt_embeds: Optional[torch.Tensor] = None,337        negative_prompt_embeds: Optional[torch.Tensor] = None,338        pooled_prompt_embeds: Optional[torch.Tensor] = None,339        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,340        ip_adapter_image: Optional[PipelineImageInput] = None,341        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,342        output_type: Optional[str] = "pil",343        return_dict: bool = True,344        cross_attention_kwargs: Optional[Dict[str, Any]] = None,345        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,346        guess_mode: bool = False,347        control_guidance_start: Union[float, List[float]] = 0.0,348        control_guidance_end: Union[float, List[float]] = 1.0,349        original_size: Tuple[int, int] = None,350        crops_coords_top_left: Tuple[int, int] = (0, 0),351        target_size: Tuple[int, int] = None,352        negative_original_size: Optional[Tuple[int, int]] = None,353        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),354        negative_target_size: Optional[Tuple[int, int]] = None,355        clip_skip: Optional[int] = None,356        callback_on_step_end: Optional[357            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]358        ] = None,359        callback_on_step_end_tensor_inputs: List[str] = ["latents"],360        attention_auto_machine_weight: float = 1.0,361        gn_auto_machine_weight: float = 1.0,362        reference_guidance_start: float = 0.0,363        reference_guidance_end: float = 1.0,364        style_fidelity: float = 0.5,365        reference_attn: bool = True,366        reference_adain: bool = True,367        **kwargs,368    ):369        r"""370        The call function to the pipeline for generation.371 372        Args:373            prompt (`str` or `List[str]`, *optional*):374                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.375            prompt_2 (`str` or `List[str]`, *optional*):376                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is377                used in both text-encoders.378            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:379                    `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):380                The ControlNet input condition to provide guidance to the `unet` for generation. If the type is381                specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted382                as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or383                width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,384                images must be passed as a list such that each element of the list can be correctly batched for input385                to a single ControlNet.386            ref_image (`torch.Tensor`, `PIL.Image.Image`):387                The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If388                the type is specified as `Torch.Tensor`, it is passed to Reference Control as is. `PIL.Image.Image` can389                also be accepted as an image.390            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):391                The height in pixels of the generated image. Anything below 512 pixels won't work well for392                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)393                and checkpoints that are not specifically fine-tuned on low resolutions.394            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):395                The width in pixels of the generated image. Anything below 512 pixels won't work well for396                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)397                and checkpoints that are not specifically fine-tuned on low resolutions.398            num_inference_steps (`int`, *optional*, defaults to 50):399                The number of denoising steps. More denoising steps usually lead to a higher quality image at the400                expense of slower inference.401            timesteps (`List[int]`, *optional*):402                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument403                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is404                passed will be used. Must be in descending order.405            sigmas (`List[float]`, *optional*):406                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in407                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed408                will be used.409            denoising_end (`float`, *optional*):410                When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be411                completed before it is intentionally prematurely terminated. As a result, the returned sample will412                still retain a substantial amount of noise as determined by the discrete timesteps selected by the413                scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a414                "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image415                Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)416            guidance_scale (`float`, *optional*, defaults to 5.0):417                A higher guidance scale value encourages the model to generate images closely linked to the text418                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.419            negative_prompt (`str` or `List[str]`, *optional*):420                The prompt or prompts to guide what to not include in image generation. If not defined, you need to421                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).422            negative_prompt_2 (`str` or `List[str]`, *optional*):423                The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2`424                and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders.425            num_images_per_prompt (`int`, *optional*, defaults to 1):426                The number of images to generate per prompt.427            eta (`float`, *optional*, defaults to 0.0):428                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies429                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.430            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):431                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make432                generation deterministic.433            latents (`torch.Tensor`, *optional*):434                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image435                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents436                tensor is generated by sampling using the supplied random `generator`.437            prompt_embeds (`torch.Tensor`, *optional*):438                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not439                provided, text embeddings are generated from the `prompt` input argument.440            negative_prompt_embeds (`torch.Tensor`, *optional*):441                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If442                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.443            pooled_prompt_embeds (`torch.Tensor`, *optional*):444                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs (prompt weighting). If445                not provided, pooled text embeddings are generated from `prompt` input argument.446            negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):447                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs (prompt448                weighting). If not provided, pooled `negative_prompt_embeds` are generated from `negative_prompt` input449                argument.450            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.451            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):452                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of453                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should454                contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not455                provided, embeddings are computed from the `ip_adapter_image` input argument.456            output_type (`str`, *optional*, defaults to `"pil"`):457                The output format of the generated image. Choose between `PIL.Image` or `np.array`.458            return_dict (`bool`, *optional*, defaults to `True`):459                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a460                plain tuple.461            cross_attention_kwargs (`dict`, *optional*):462                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in463                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).464            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):465                The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added466                to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set467                the corresponding scale as a list.468            guess_mode (`bool`, *optional*, defaults to `False`):469                The ControlNet encoder tries to recognize the content of the input image even if you remove all470                prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.471            control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):472                The percentage of total steps at which the ControlNet starts applying.473            control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):474                The percentage of total steps at which the ControlNet stops applying.475            original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):476                If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.477                `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as478                explained in section 2.2 of479                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).480            crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):481                `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position482                `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting483                `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of484                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).485            target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):486                For most cases, `target_size` should be set to the desired height and width of the generated image. If487                not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in488                section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).489            negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):490                To negatively condition the generation process based on a specific image resolution. Part of SDXL's491                micro-conditioning as explained in section 2.2 of492                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more493                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.494            negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):495                To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's496                micro-conditioning as explained in section 2.2 of497                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more498                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.499            negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):500                To negatively condition the generation process based on a target image resolution. It should be as same501                as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of502                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more503                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.504            clip_skip (`int`, *optional*):505                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that506                the output of the pre-final layer will be used for computing the prompt embeddings.507            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):508                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of509                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:510                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a511                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.512            callback_on_step_end_tensor_inputs (`List`, *optional*):513                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list514                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the515                `._callback_tensor_inputs` attribute of your pipeline class.516            attention_auto_machine_weight (`float`):517                Weight of using reference query for self attention's context.518                If attention_auto_machine_weight=1.0, use reference query for all self attention's context.519            gn_auto_machine_weight (`float`):520                Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.521            reference_guidance_start (`float`, *optional*, defaults to 0.0):522                The percentage of total steps at which the reference ControlNet starts applying.523            reference_guidance_end (`float`, *optional*, defaults to 1.0):524                The percentage of total steps at which the reference ControlNet stops applying.525            style_fidelity (`float`):526                style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,527                elif style_fidelity=0.0, prompt more important, else balanced.528            reference_attn (`bool`):529                Whether to use reference query for self attention's context.530            reference_adain (`bool`):531                Whether to use reference adain.532 533        Examples:534 535        Returns:536            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:537                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,538                otherwise a `tuple` is returned containing the output images.539        """540 541        callback = kwargs.pop("callback", None)542        callback_steps = kwargs.pop("callback_steps", None)543 544        if callback is not None:545            deprecate(546                "callback",547                "1.0.0",548                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",549            )550        if callback_steps is not None:551            deprecate(552                "callback_steps",553                "1.0.0",554                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",555            )556 557        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):558            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs559 560        controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet561 562        # align format for control guidance563        if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):564            control_guidance_start = len(control_guidance_end) * [control_guidance_start]565        elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):566            control_guidance_end = len(control_guidance_start) * [control_guidance_end]567        elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):568            mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1569            control_guidance_start, control_guidance_end = (570                mult * [control_guidance_start],571                mult * [control_guidance_end],572            )573 574        # 1. Check inputs. Raise error if not correct575        self.check_inputs(576            prompt,577            prompt_2,578            image,579            callback_steps,580            negative_prompt,581            negative_prompt_2,582            prompt_embeds,583            negative_prompt_embeds,584            pooled_prompt_embeds,585            ip_adapter_image,586            ip_adapter_image_embeds,587            negative_pooled_prompt_embeds,588            controlnet_conditioning_scale,589            control_guidance_start,590            control_guidance_end,591            callback_on_step_end_tensor_inputs,592        )593 594        self.check_ref_inputs(595            ref_image,596            reference_guidance_start,597            reference_guidance_end,598            style_fidelity,599            reference_attn,600            reference_adain,601        )602 603        self._guidance_scale = guidance_scale604        self._clip_skip = clip_skip605        self._cross_attention_kwargs = cross_attention_kwargs606        self._denoising_end = denoising_end607        self._interrupt = False608 609        # 2. Define call parameters610        if prompt is not None and isinstance(prompt, str):611            batch_size = 1612        elif prompt is not None and isinstance(prompt, list):613            batch_size = len(prompt)614        else:615            batch_size = prompt_embeds.shape[0]616 617        device = self._execution_device618 619        if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):620            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)621 622        global_pool_conditions = (623            controlnet.config.global_pool_conditions624            if isinstance(controlnet, ControlNetModel)625            else controlnet.nets[0].config.global_pool_conditions626        )627        guess_mode = guess_mode or global_pool_conditions628 629        # 3.1 Encode input prompt630        text_encoder_lora_scale = (631            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None632        )633        (634            prompt_embeds,635            negative_prompt_embeds,636            pooled_prompt_embeds,637            negative_pooled_prompt_embeds,638        ) = self.encode_prompt(639            prompt,640            prompt_2,641            device,642            num_images_per_prompt,643            self.do_classifier_free_guidance,644            negative_prompt,645            negative_prompt_2,646            prompt_embeds=prompt_embeds,647            negative_prompt_embeds=negative_prompt_embeds,648            pooled_prompt_embeds=pooled_prompt_embeds,649            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,650            lora_scale=text_encoder_lora_scale,651            clip_skip=self.clip_skip,652        )653 654        # 3.2 Encode ip_adapter_image655        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:656            image_embeds = self.prepare_ip_adapter_image_embeds(657                ip_adapter_image,658                ip_adapter_image_embeds,659                device,660                batch_size * num_images_per_prompt,661                self.do_classifier_free_guidance,662            )663 664        # 4. Prepare image665        if isinstance(controlnet, ControlNetModel):666            image = self.prepare_image(667                image=image,668                width=width,669                height=height,670                batch_size=batch_size * num_images_per_prompt,671                num_images_per_prompt=num_images_per_prompt,672                device=device,673                dtype=controlnet.dtype,674                do_classifier_free_guidance=self.do_classifier_free_guidance,675                guess_mode=guess_mode,676            )677            height, width = image.shape[-2:]678        elif isinstance(controlnet, MultiControlNetModel):679            images = []680 681            for image_ in image:682                image_ = self.prepare_image(683                    image=image_,684                    width=width,685                    height=height,686                    batch_size=batch_size * num_images_per_prompt,687                    num_images_per_prompt=num_images_per_prompt,688                    device=device,689                    dtype=controlnet.dtype,690                    do_classifier_free_guidance=self.do_classifier_free_guidance,691                    guess_mode=guess_mode,692                )693 694                images.append(image_)695 696            image = images697            height, width = image[0].shape[-2:]698        else:699            assert False700 701        # 5. Preprocess reference image702        ref_image = self.prepare_ref_image(703            image=ref_image,704            width=width,705            height=height,706            batch_size=batch_size * num_images_per_prompt,707            num_images_per_prompt=num_images_per_prompt,708            device=device,709            dtype=prompt_embeds.dtype,710        )711 712        # 6. Prepare timesteps713        timesteps, num_inference_steps = retrieve_timesteps(714            self.scheduler, num_inference_steps, device, timesteps, sigmas715        )716        self._num_timesteps = len(timesteps)717 718        # 7. Prepare latent variables719        num_channels_latents = self.unet.config.in_channels720        latents = self.prepare_latents(721            batch_size * num_images_per_prompt,722            num_channels_latents,723            height,724            width,725            prompt_embeds.dtype,726            device,727            generator,728            latents,729        )730 731        # 7.5 Optionally get Guidance Scale Embedding732        timestep_cond = None733        if self.unet.config.time_cond_proj_dim is not None:734            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)735            timestep_cond = self.get_guidance_scale_embedding(736                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim737            ).to(device=device, dtype=latents.dtype)738 739        # 8. Prepare reference latent variables740        ref_image_latents = self.prepare_ref_latents(741            ref_image,742            batch_size * num_images_per_prompt,743            prompt_embeds.dtype,744            device,745            generator,746            self.do_classifier_free_guidance,747        )748 749        # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline750        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)751 752        # 9.1 Create tensor stating which controlnets to keep753        controlnet_keep = []754        reference_keeps = []755        for i in range(len(timesteps)):756            keeps = [757                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)758                for s, e in zip(control_guidance_start, control_guidance_end)759            ]760            controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps)761            reference_keep = 1.0 - float(762                i / len(timesteps) < reference_guidance_start or (i + 1) / len(timesteps) > reference_guidance_end763            )764            reference_keeps.append(reference_keep)765 766        # 9.2 Modify self attention and group norm767        MODE = "write"768        uc_mask = (769            torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)770            .type_as(ref_image_latents)771            .bool()772        )773 774        do_classifier_free_guidance = self.do_classifier_free_guidance775 776        def hacked_basic_transformer_inner_forward(777            self,778            hidden_states: torch.Tensor,779            attention_mask: Optional[torch.Tensor] = None,780            encoder_hidden_states: Optional[torch.Tensor] = None,781            encoder_attention_mask: Optional[torch.Tensor] = None,782            timestep: Optional[torch.LongTensor] = None,783            cross_attention_kwargs: Dict[str, Any] = None,784            class_labels: Optional[torch.LongTensor] = None,785        ):786            if self.use_ada_layer_norm:787                norm_hidden_states = self.norm1(hidden_states, timestep)788            elif self.use_ada_layer_norm_zero:789                norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(790                    hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype791                )792            else:793                norm_hidden_states = self.norm1(hidden_states)794 795            # 1. Self-Attention796            cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}797            if self.only_cross_attention:798                attn_output = self.attn1(799                    norm_hidden_states,800                    encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,801                    attention_mask=attention_mask,802                    **cross_attention_kwargs,803                )804            else:805                if MODE == "write":806                    self.bank.append(norm_hidden_states.detach().clone())807                    attn_output = self.attn1(808                        norm_hidden_states,809                        encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,810                        attention_mask=attention_mask,811                        **cross_attention_kwargs,812                    )813                if MODE == "read":814                    if attention_auto_machine_weight > self.attn_weight:815                        attn_output_uc = self.attn1(816                            norm_hidden_states,817                            encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),818                            # attention_mask=attention_mask,819                            **cross_attention_kwargs,820                        )821                        attn_output_c = attn_output_uc.clone()822                        if do_classifier_free_guidance and style_fidelity > 0:823                            attn_output_c[uc_mask] = self.attn1(824                                norm_hidden_states[uc_mask],825                                encoder_hidden_states=norm_hidden_states[uc_mask],826                                **cross_attention_kwargs,827                            )828                        attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc829                        self.bank.clear()830                    else:831                        attn_output = self.attn1(832                            norm_hidden_states,833                            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,834                            attention_mask=attention_mask,835                            **cross_attention_kwargs,836                        )837            if self.use_ada_layer_norm_zero:838                attn_output = gate_msa.unsqueeze(1) * attn_output839            hidden_states = attn_output + hidden_states840 841            if self.attn2 is not None:842                norm_hidden_states = (843                    self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)844                )845 846                # 2. Cross-Attention847                attn_output = self.attn2(848                    norm_hidden_states,849                    encoder_hidden_states=encoder_hidden_states,850                    attention_mask=encoder_attention_mask,851                    **cross_attention_kwargs,852                )853                hidden_states = attn_output + hidden_states854 855            # 3. Feed-forward856            norm_hidden_states = self.norm3(hidden_states)857 858            if self.use_ada_layer_norm_zero:859                norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]860 861            ff_output = self.ff(norm_hidden_states)862 863            if self.use_ada_layer_norm_zero:864                ff_output = gate_mlp.unsqueeze(1) * ff_output865 866            hidden_states = ff_output + hidden_states867 868            return hidden_states869 870        def hacked_mid_forward(self, *args, **kwargs):871            eps = 1e-6872            x = self.original_forward(*args, **kwargs)873            if MODE == "write":874                if gn_auto_machine_weight >= self.gn_weight:875                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)876                    self.mean_bank.append(mean)877                    self.var_bank.append(var)878            if MODE == "read":879                if len(self.mean_bank) > 0 and len(self.var_bank) > 0:880                    var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)881                    std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5882                    mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))883                    var_acc = sum(self.var_bank) / float(len(self.var_bank))884                    std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5885                    x_uc = (((x - mean) / std) * std_acc) + mean_acc886                    x_c = x_uc.clone()887                    if do_classifier_free_guidance and style_fidelity > 0:888                        x_c[uc_mask] = x[uc_mask]889                    x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc890                self.mean_bank = []891                self.var_bank = []892            return x893 894        def hack_CrossAttnDownBlock2D_forward(895            self,896            hidden_states: torch.Tensor,897            temb: Optional[torch.Tensor] = None,898            encoder_hidden_states: Optional[torch.Tensor] = None,899            attention_mask: Optional[torch.Tensor] = None,900            cross_attention_kwargs: Optional[Dict[str, Any]] = None,901            encoder_attention_mask: Optional[torch.Tensor] = None,902        ):903            eps = 1e-6904 905            # TODO(Patrick, William) - attention mask is not used906            output_states = ()907 908            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):909                hidden_states = resnet(hidden_states, temb)910                hidden_states = attn(911                    hidden_states,912                    encoder_hidden_states=encoder_hidden_states,913                    cross_attention_kwargs=cross_attention_kwargs,914                    attention_mask=attention_mask,915                    encoder_attention_mask=encoder_attention_mask,916                    return_dict=False,917                )[0]918                if MODE == "write":919                    if gn_auto_machine_weight >= self.gn_weight:920                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)921                        self.mean_bank.append([mean])922                        self.var_bank.append([var])923                if MODE == "read":924                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:925                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)926                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5927                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))928                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))929                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5930                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc931                        hidden_states_c = hidden_states_uc.clone()932                        if do_classifier_free_guidance and style_fidelity > 0:933                            hidden_states_c[uc_mask] = hidden_states[uc_mask]934                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc935 936                output_states = output_states + (hidden_states,)937 938            if MODE == "read":939                self.mean_bank = []940                self.var_bank = []941 942            if self.downsamplers is not None:943                for downsampler in self.downsamplers:944                    hidden_states = downsampler(hidden_states)945 946                output_states = output_states + (hidden_states,)947 948            return hidden_states, output_states949 950        def hacked_DownBlock2D_forward(self, hidden_states, temb=None, *args, **kwargs):951            eps = 1e-6952 953            output_states = ()954 955            for i, resnet in enumerate(self.resnets):956                hidden_states = resnet(hidden_states, temb)957 958                if MODE == "write":959                    if gn_auto_machine_weight >= self.gn_weight:960                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)961                        self.mean_bank.append([mean])962                        self.var_bank.append([var])963                if MODE == "read":964                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:965                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)966                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5967                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))968                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))969                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5970                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc971                        hidden_states_c = hidden_states_uc.clone()972                        if do_classifier_free_guidance and style_fidelity > 0:973                            hidden_states_c[uc_mask] = hidden_states[uc_mask]974                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc975 976                output_states = output_states + (hidden_states,)977 978            if MODE == "read":979                self.mean_bank = []980                self.var_bank = []981 982            if self.downsamplers is not None:983                for downsampler in self.downsamplers:984                    hidden_states = downsampler(hidden_states)985 986                output_states = output_states + (hidden_states,)987 988            return hidden_states, output_states989 990        def hacked_CrossAttnUpBlock2D_forward(991            self,992            hidden_states: torch.Tensor,993            res_hidden_states_tuple: Tuple[torch.Tensor, ...],994            temb: Optional[torch.Tensor] = None,995            encoder_hidden_states: Optional[torch.Tensor] = None,996            cross_attention_kwargs: Optional[Dict[str, Any]] = None,997            upsample_size: Optional[int] = None,998            attention_mask: Optional[torch.Tensor] = None,999            encoder_attention_mask: Optional[torch.Tensor] = None,1000        ):1001            eps = 1e-61002            # TODO(Patrick, William) - attention mask is not used1003            for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):1004                # pop res hidden states1005                res_hidden_states = res_hidden_states_tuple[-1]1006                res_hidden_states_tuple = res_hidden_states_tuple[:-1]1007                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)1008                hidden_states = resnet(hidden_states, temb)1009                hidden_states = attn(1010                    hidden_states,1011                    encoder_hidden_states=encoder_hidden_states,1012                    cross_attention_kwargs=cross_attention_kwargs,1013                    attention_mask=attention_mask,1014                    encoder_attention_mask=encoder_attention_mask,1015                    return_dict=False,1016                )[0]1017 1018                if MODE == "write":1019                    if gn_auto_machine_weight >= self.gn_weight:1020                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1021                        self.mean_bank.append([mean])1022                        self.var_bank.append([var])1023                if MODE == "read":1024                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1025                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1026                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51027                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))1028                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))1029                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51030                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc1031                        hidden_states_c = hidden_states_uc.clone()1032                        if do_classifier_free_guidance and style_fidelity > 0:1033                            hidden_states_c[uc_mask] = hidden_states[uc_mask]1034                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc1035 1036            if MODE == "read":1037                self.mean_bank = []1038                self.var_bank = []1039 1040            if self.upsamplers is not None:1041                for upsampler in self.upsamplers:1042                    hidden_states = upsampler(hidden_states, upsample_size)1043 1044            return hidden_states1045 1046        def hacked_UpBlock2D_forward(1047            self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, *args, **kwargs1048        ):1049            eps = 1e-61050            for i, resnet in enumerate(self.resnets):1051                # pop res hidden states1052                res_hidden_states = res_hidden_states_tuple[-1]1053                res_hidden_states_tuple = res_hidden_states_tuple[:-1]1054                hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)1055                hidden_states = resnet(hidden_states, temb)1056 1057                if MODE == "write":1058                    if gn_auto_machine_weight >= self.gn_weight:1059                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1060                        self.mean_bank.append([mean])1061                        self.var_bank.append([var])1062                if MODE == "read":1063                    if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1064                        var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1065                        std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51066                        mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))1067                        var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))1068                        std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51069                        hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc1070                        hidden_states_c = hidden_states_uc.clone()1071                        if do_classifier_free_guidance and style_fidelity > 0:1072                            hidden_states_c[uc_mask] = hidden_states[uc_mask]1073                        hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc1074 1075            if MODE == "read":1076                self.mean_bank = []1077                self.var_bank = []1078 1079            if self.upsamplers is not None:1080                for upsampler in self.upsamplers:1081                    hidden_states = upsampler(hidden_states, upsample_size)1082 1083            return hidden_states1084 1085        if reference_attn:1086            attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]1087            attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])1088 1089            for i, module in enumerate(attn_modules):1090                module._original_inner_forward = module.forward1091                module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)1092                module.bank = []1093                module.attn_weight = float(i) / float(len(attn_modules))1094 1095        if reference_adain:1096            gn_modules = [self.unet.mid_block]1097            self.unet.mid_block.gn_weight = 01098 1099            down_blocks = self.unet.down_blocks1100            for w, module in enumerate(down_blocks):1101                module.gn_weight = 1.0 - float(w) / float(len(down_blocks))1102                gn_modules.append(module)1103 1104            up_blocks = self.unet.up_blocks1105            for w, module in enumerate(up_blocks):1106                module.gn_weight = float(w) / float(len(up_blocks))1107                gn_modules.append(module)1108 1109            for i, module in enumerate(gn_modules):1110                if getattr(module, "original_forward", None) is None:1111                    module.original_forward = module.forward1112                if i == 0:1113                    # mid_block1114                    module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)1115                elif isinstance(module, CrossAttnDownBlock2D):1116                    module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)1117                elif isinstance(module, DownBlock2D):1118                    module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)1119                elif isinstance(module, CrossAttnUpBlock2D):1120                    module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)1121                elif isinstance(module, UpBlock2D):1122                    module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)1123                module.mean_bank = []1124                module.var_bank = []1125                module.gn_weight *= 21126 1127        # 9.2 Prepare added time ids & embeddings1128        if isinstance(image, list):1129            original_size = original_size or image[0].shape[-2:]1130        else:1131            original_size = original_size or image.shape[-2:]1132        target_size = target_size or (height, width)1133 1134        add_text_embeds = pooled_prompt_embeds1135        if self.text_encoder_2 is None:1136            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])1137        else:1138            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim1139 1140        add_time_ids = self._get_add_time_ids(1141            original_size,1142            crops_coords_top_left,1143            target_size,1144            dtype=prompt_embeds.dtype,1145            text_encoder_projection_dim=text_encoder_projection_dim,1146        )1147 1148        if negative_original_size is not None and negative_target_size is not None:1149            negative_add_time_ids = self._get_add_time_ids(1150                negative_original_size,1151                negative_crops_coords_top_left,1152                negative_target_size,1153                dtype=prompt_embeds.dtype,1154                text_encoder_projection_dim=text_encoder_projection_dim,1155            )1156        else:1157            negative_add_time_ids = add_time_ids1158 1159        if self.do_classifier_free_guidance:1160            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)1161            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)1162            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)1163 1164        prompt_embeds = prompt_embeds.to(device)1165        add_text_embeds = add_text_embeds.to(device)1166        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)1167 1168        # 10. Denoising loop1169        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order1170 1171        # 10.1 Apply denoising_end1172        if (1173            self.denoising_end is not None1174            and isinstance(self.denoising_end, float)1175            and self.denoising_end > 01176            and self.denoising_end < 11177        ):1178            discrete_timestep_cutoff = int(1179                round(1180                    self.scheduler.config.num_train_timesteps1181                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)1182                )1183            )1184            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))1185            timesteps = timesteps[:num_inference_steps]1186 1187        is_unet_compiled = is_compiled_module(self.unet)1188        is_controlnet_compiled = is_compiled_module(self.controlnet)1189        is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1")1190        with self.progress_bar(total=num_inference_steps) as progress_bar:1191            for i, t in enumerate(timesteps):1192                if self.interrupt:1193                    continue1194 1195                # Relevant thread:1196                # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/14281197                if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1:1198                    torch._inductor.cudagraph_mark_step_begin()1199                # expand the latents if we are doing classifier free guidance1200                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents

Showing the first 1,200 of 1363 lines. Download the file for the rest.