Delete pipeline_fill_sd_xl.py

pipeline_fill_sd_xl.py

@@ -1,559 +0,0 @@
-# Copyright 2024 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.
-
-from typing import List, Optional, Union
-
-import cv2
-import PIL.Image
-import torch
-import torch.nn.functional as F
-from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
-from diffusers.models import AutoencoderKL, UNet2DConditionModel
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
-from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils.torch_utils import randn_tensor
-from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
-
-from controlnet_union import ControlNetModel_Union
-
-
-def latents_to_rgb(latents):
-    weights = ((60, -60, 25, -70), (60, -5, 15, -50), (60, 10, -5, -35))
-
-    weights_tensor = torch.t(
-        torch.tensor(weights, dtype=latents.dtype).to(latents.device)
-    )
-    biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(
-        latents.device
-    )
-    rgb_tensor = torch.einsum(
-        "...lxy,lr -> ...rxy", latents, weights_tensor
-    ) + biases_tensor.unsqueeze(-1).unsqueeze(-1)
-    image_array = rgb_tensor.clamp(0, 255)[0].byte().cpu().numpy()
-    image_array = image_array.transpose(1, 2, 0)  # Change the order of dimensions
-
-    denoised_image = cv2.fastNlMeansDenoisingColored(image_array, None, 10, 10, 7, 21)
-    blurred_image = cv2.GaussianBlur(denoised_image, (5, 5), 0)
-    final_image = PIL.Image.fromarray(blurred_image)
-
-    width, height = final_image.size
-    final_image = final_image.resize(
-        (width * 8, height * 8), PIL.Image.Resampling.LANCZOS
-    )
-
-    return final_image
-
-
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    **kwargs,
-):
-    scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-    timesteps = scheduler.timesteps
-
-    return timesteps, num_inference_steps
-
-
-class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin):
-    model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
-    _optional_components = [
-        "tokenizer",
-        "tokenizer_2",
-        "text_encoder",
-        "text_encoder_2",
-    ]
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: CLIPTextModel,
-        text_encoder_2: CLIPTextModelWithProjection,
-        tokenizer: CLIPTokenizer,
-        tokenizer_2: CLIPTokenizer,
-        unet: UNet2DConditionModel,
-        controlnet: ControlNetModel_Union,
-        scheduler: KarrasDiffusionSchedulers,
-        force_zeros_for_empty_prompt: bool = True,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            text_encoder_2=text_encoder_2,
-            tokenizer=tokenizer,
-            tokenizer_2=tokenizer_2,
-            unet=unet,
-            controlnet=controlnet,
-            scheduler=scheduler,
-        )
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.image_processor = VaeImageProcessor(
-            vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True
-        )
-        self.control_image_processor = VaeImageProcessor(
-            vae_scale_factor=self.vae_scale_factor,
-            do_convert_rgb=True,
-            do_normalize=False,
-        )
-
-        self.register_to_config(
-            force_zeros_for_empty_prompt=force_zeros_for_empty_prompt
-        )
-
-    def encode_prompt(
-        self,
-        prompt: str,
-        device: Optional[torch.device] = None,
-        do_classifier_free_guidance: bool = True,
-    ):
-        device = device or self._execution_device
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        if prompt is not None:
-            batch_size = len(prompt)
-
-        # Define tokenizers and text encoders
-        tokenizers = (
-            [self.tokenizer, self.tokenizer_2]
-            if self.tokenizer is not None
-            else [self.tokenizer_2]
-        )
-        text_encoders = (
-            [self.text_encoder, self.text_encoder_2]
-            if self.text_encoder is not None
-            else [self.text_encoder_2]
-        )
-
-        prompt_2 = prompt
-        prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
-
-        # textual inversion: process multi-vector tokens if necessary
-        prompt_embeds_list = []
-        prompts = [prompt, prompt_2]
-        for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
-            text_inputs = tokenizer(
-                prompt,
-                padding="max_length",
-                max_length=tokenizer.model_max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-
-            text_input_ids = text_inputs.input_ids
-
-            prompt_embeds = text_encoder(
-                text_input_ids.to(device), output_hidden_states=True
-            )
-
-            # We are only ALWAYS interested in the pooled output of the final text encoder
-            pooled_prompt_embeds = prompt_embeds[0]
-            prompt_embeds = prompt_embeds.hidden_states[-2]
-            prompt_embeds_list.append(prompt_embeds)
-
-        prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-
-        # get unconditional embeddings for classifier free guidance
-        zero_out_negative_prompt = True
-        negative_prompt_embeds = None
-        negative_pooled_prompt_embeds = None
-
-        if do_classifier_free_guidance and zero_out_negative_prompt:
-            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
-            negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
-        elif do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = ""
-            negative_prompt_2 = 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
-            )
-
-            uncond_tokens: List[str]
-            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`."
-                )
-            else:
-                uncond_tokens = [negative_prompt, negative_prompt_2]
-
-            negative_prompt_embeds_list = []
-            for negative_prompt, tokenizer, text_encoder in zip(
-                uncond_tokens, tokenizers, text_encoders
-            ):
-                max_length = prompt_embeds.shape[1]
-                uncond_input = tokenizer(
-                    negative_prompt,
-                    padding="max_length",
-                    max_length=max_length,
-                    truncation=True,
-                    return_tensors="pt",
-                )
-
-                negative_prompt_embeds = text_encoder(
-                    uncond_input.input_ids.to(device),
-                    output_hidden_states=True,
-                )
-                # We are only ALWAYS interested in the pooled output of the final text encoder
-                negative_pooled_prompt_embeds = negative_prompt_embeds[0]
-                negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
-
-                negative_prompt_embeds_list.append(negative_prompt_embeds)
-
-            negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
-
-        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
-
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, 1, 1)
-        prompt_embeds = prompt_embeds.view(bs_embed * 1, seq_len, -1)
-
-        if do_classifier_free_guidance:
-            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
-            seq_len = negative_prompt_embeds.shape[1]
-
-            if self.text_encoder_2 is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(
-                    dtype=self.text_encoder_2.dtype, device=device
-                )
-            else:
-                negative_prompt_embeds = negative_prompt_embeds.to(
-                    dtype=self.unet.dtype, device=device
-                )
-
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(
-                batch_size * 1, seq_len, -1
-            )
-
-        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1)
-        if do_classifier_free_guidance:
-            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(
-                1, 1
-            ).view(bs_embed * 1, -1)
-
-        return (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        )
-
-    def check_inputs(
-        self,
-        prompt_embeds,
-        negative_prompt_embeds,
-        pooled_prompt_embeds,
-        negative_pooled_prompt_embeds,
-        image,
-        controlnet_conditioning_scale=1.0,
-    ):
-        if prompt_embeds is None:
-            raise ValueError(
-                "Provide `prompt_embeds`. Cannot leave `prompt_embeds` undefined."
-            )
-
-        if negative_prompt_embeds is None:
-            raise ValueError(
-                "Provide `negative_prompt_embeds`. Cannot leave `negative_prompt_embeds` undefined."
-            )
-
-        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`."
-            )
-
-        # Check `image`
-        is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
-            self.controlnet, torch._dynamo.eval_frame.OptimizedModule
-        )
-        if (
-            isinstance(self.controlnet, ControlNetModel_Union)
-            or is_compiled
-            and isinstance(self.controlnet._orig_mod, ControlNetModel_Union)
-        ):
-            if not isinstance(image, PIL.Image.Image):
-                raise TypeError(
-                    f"image must be passed and has to be a PIL image, but is {type(image)}"
-                )
-
-        else:
-            assert False
-
-        # Check `controlnet_conditioning_scale`
-        if (
-            isinstance(self.controlnet, ControlNetModel_Union)
-            or is_compiled
-            and isinstance(self.controlnet._orig_mod, ControlNetModel_Union)
-        ):
-            if not isinstance(controlnet_conditioning_scale, float):
-                raise TypeError(
-                    "For single controlnet: `controlnet_conditioning_scale` must be type `float`."
-                )
-        else:
-            assert False
-
-    def prepare_image(self, image, device, dtype, do_classifier_free_guidance=False):
-        image = self.control_image_processor.preprocess(image).to(dtype=torch.float32)
-
-        image_batch_size = image.shape[0]
-
-        image = image.repeat_interleave(image_batch_size, dim=0)
-        image = image.to(device=device, dtype=dtype)
-
-        if do_classifier_free_guidance:
-            image = torch.cat([image] * 2)
-
-        return image
-
-    def prepare_latents(
-        self, batch_size, num_channels_latents, height, width, dtype, device
-    ):
-        shape = (
-            batch_size,
-            num_channels_latents,
-            int(height) // self.vae_scale_factor,
-            int(width) // self.vae_scale_factor,
-        )
-
-        latents = randn_tensor(shape, device=device, dtype=dtype)
-
-        # scale the initial noise by the standard deviation required by the scheduler
-        latents = latents * self.scheduler.init_noise_sigma
-        return latents
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    # 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 and self.unet.config.time_cond_proj_dim is None
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt_embeds: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor,
-        pooled_prompt_embeds: torch.Tensor,
-        negative_pooled_prompt_embeds: torch.Tensor,
-        image: PipelineImageInput = None,
-        num_inference_steps: int = 8,
-        guidance_scale: float = 1.5,
-        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
-    ):
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-            image,
-            controlnet_conditioning_scale,
-        )
-
-        self._guidance_scale = guidance_scale
-
-        # 2. Define call parameters
-        batch_size = 1
-        device = self._execution_device
-
-        # 4. Prepare image
-        if isinstance(self.controlnet, ControlNetModel_Union):
-            image = self.prepare_image(
-                image=image,
-                device=device,
-                dtype=self.controlnet.dtype,
-                do_classifier_free_guidance=self.do_classifier_free_guidance,
-            )
-            height, width = image.shape[-2:]
-        else:
-            assert False
-
-        # 5. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler, num_inference_steps, device
-        )
-        self._num_timesteps = len(timesteps)
-
-        # 6. Prepare latent variables
-        num_channels_latents = self.unet.config.in_channels
-        latents = self.prepare_latents(
-            batch_size,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-        )
-
-        # 7 Prepare added time ids & embeddings
-        add_text_embeds = pooled_prompt_embeds
-
-        add_time_ids = negative_add_time_ids = torch.tensor(
-            image.shape[-2:] + torch.Size([0, 0]) + image.shape[-2:]
-        ).unsqueeze(0)
-
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            add_text_embeds = torch.cat(
-                [negative_pooled_prompt_embeds, add_text_embeds], dim=0
-            )
-            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
-
-        prompt_embeds = prompt_embeds.to(device)
-        add_text_embeds = add_text_embeds.to(device)
-        add_time_ids = add_time_ids.to(device).repeat(batch_size, 1)
-
-        controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0]
-        union_control_type = (
-            torch.Tensor([0, 0, 0, 0, 0, 0, 1, 0])
-            .to(device, dtype=prompt_embeds.dtype)
-            .repeat(batch_size * 2, 1)
-        )
-
-        added_cond_kwargs = {
-            "text_embeds": add_text_embeds,
-            "time_ids": add_time_ids,
-            "control_type": union_control_type,
-        }
-
-        controlnet_prompt_embeds = prompt_embeds
-        controlnet_added_cond_kwargs = added_cond_kwargs
-
-        # 8. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # 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
-                )
-                latent_model_input = self.scheduler.scale_model_input(
-                    latent_model_input, t
-                )
-
-                # controlnet(s) inference
-                control_model_input = latent_model_input
-
-                down_block_res_samples, mid_block_res_sample = self.controlnet(
-                    control_model_input,
-                    t,
-                    encoder_hidden_states=controlnet_prompt_embeds,
-                    controlnet_cond_list=controlnet_image_list,
-                    conditioning_scale=controlnet_conditioning_scale,
-                    guess_mode=False,
-                    added_cond_kwargs=controlnet_added_cond_kwargs,
-                    return_dict=False,
-                )
-
-                # predict the noise residual
-                noise_pred = self.unet(
-                    latent_model_input,
-                    t,
-                    encoder_hidden_states=prompt_embeds,
-                    timestep_cond=None,
-                    cross_attention_kwargs={},
-                    down_block_additional_residuals=down_block_res_samples,
-                    mid_block_additional_residual=mid_block_res_sample,
-                    added_cond_kwargs=added_cond_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 + guidance_scale * (
-                        noise_pred_text - noise_pred_uncond
-                    )
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(
-                    noise_pred, t, latents, return_dict=False
-                )[0]
-
-                if i == 2:
-                    prompt_embeds = prompt_embeds[-1:]
-                    add_text_embeds = add_text_embeds[-1:]
-                    add_time_ids = add_time_ids[-1:]
-                    union_control_type = union_control_type[-1:]
-
-                    added_cond_kwargs = {
-                        "text_embeds": add_text_embeds,
-                        "time_ids": add_time_ids,
-                        "control_type": union_control_type,
-                    }
-
-                    controlnet_prompt_embeds = prompt_embeds
-                    controlnet_added_cond_kwargs = added_cond_kwargs
-
-                    image = image[-1:]
-                    controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0]
-
-                    self._guidance_scale = 0.0
-
-                if i == len(timesteps) - 1 or (
-                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
-                ):
-                    progress_bar.update()
-                    yield latents_to_rgb(latents)
-
-        latents = latents / self.vae.config.scaling_factor
-        image = self.vae.decode(latents, return_dict=False)[0]
-        image = self.image_processor.postprocess(image)[0]
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        yield image