multimodalart/pix2pix-zero

1 2import torch3import inspect4from packaging import version5from typing import Any, Callable, Dict, List, Optional, Union6 7from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer8from diffusers import DiffusionPipeline9from diffusers.models import AutoencoderKL, UNet2DConditionModel10from diffusers.schedulers import KarrasDiffusionSchedulers11from diffusers.utils import deprecate, is_accelerate_available, logging, randn_tensor, replace_example_docstring12from diffusers import StableDiffusionPipeline13from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker14 15 16 17class BasePipeline(DiffusionPipeline):18    _optional_components = ["safety_checker", "feature_extractor"]19    def __init__(20        self,21        vae: AutoencoderKL,22        text_encoder: CLIPTextModel,23        tokenizer: CLIPTokenizer,24        unet: UNet2DConditionModel,25        scheduler: KarrasDiffusionSchedulers,26        safety_checker: StableDiffusionSafetyChecker,27        feature_extractor: CLIPFeatureExtractor,28        requires_safety_checker: bool = True,29    ):30        super().__init__()31 32        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:33            deprecation_message = (34                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"35                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "36                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"37                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"38                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"39                " file"40            )41            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)42            new_config = dict(scheduler.config)43            new_config["steps_offset"] = 144            scheduler._internal_dict = FrozenDict(new_config)45 46        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:47            deprecation_message = (48                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."49                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"50                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"51                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"52                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"53            )54            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)55            new_config = dict(scheduler.config)56            new_config["clip_sample"] = False57            scheduler._internal_dict = FrozenDict(new_config)58 59        if safety_checker is None and requires_safety_checker:60            logger.warning(61                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"62                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"63                " results in services or applications open to the public. Both the diffusers team and Hugging Face"64                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"65                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"66                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."67            )68 69        if safety_checker is not None and feature_extractor is None:70            raise ValueError(71                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"72                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."73            )74 75        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(76            version.parse(unet.config._diffusers_version).base_version77        ) < version.parse("0.9.0.dev0")78        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 6479        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:80            deprecation_message = (81                "The configuration file of the unet has set the default `sample_size` to smaller than"82                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"83                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"84                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"85                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"86                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"87                " in the config might lead to incorrect results in future versions. If you have downloaded this"88                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"89                " the `unet/config.json` file"90            )91            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)92            new_config = dict(unet.config)93            new_config["sample_size"] = 6494            unet._internal_dict = FrozenDict(new_config)95 96        self.register_modules(97            vae=vae,98            text_encoder=text_encoder,99            tokenizer=tokenizer,100            unet=unet,101            scheduler=scheduler,102            safety_checker=safety_checker,103            feature_extractor=feature_extractor,104        )105        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)106        self.register_to_config(requires_safety_checker=requires_safety_checker)107 108    @property109    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device110    def _execution_device(self):111        r"""112        Returns the device on which the pipeline's models will be executed. After calling113        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module114        hooks.115        """116        if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):117            return self.device118        for module in self.unet.modules():119            if (120                hasattr(module, "_hf_hook")121                and hasattr(module._hf_hook, "execution_device")122                and module._hf_hook.execution_device is not None123            ):124                return torch.device(module._hf_hook.execution_device)125        return self.device126 127 128    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt129    def _encode_prompt(130        self,131        prompt,132        device,133        num_images_per_prompt,134        do_classifier_free_guidance,135        negative_prompt=None,136        prompt_embeds: Optional[torch.FloatTensor] = None,137        negative_prompt_embeds: Optional[torch.FloatTensor] = None,138    ):139        r"""140        Encodes the prompt into text encoder hidden states.141 142        Args:143             prompt (`str` or `List[str]`, *optional*):144                prompt to be encoded145            device: (`torch.device`):146                torch device147            num_images_per_prompt (`int`):148                number of images that should be generated per prompt149            do_classifier_free_guidance (`bool`):150                whether to use classifier free guidance or not151            negative_ prompt (`str` or `List[str]`, *optional*):152                The prompt or prompts not to guide the image generation. If not defined, one has to pass153                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.154                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).155            prompt_embeds (`torch.FloatTensor`, *optional*):156                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not157                provided, text embeddings will be generated from `prompt` input argument.158            negative_prompt_embeds (`torch.FloatTensor`, *optional*):159                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt160                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input161                argument.162        """163        if prompt is not None and isinstance(prompt, str):164            batch_size = 1165        elif prompt is not None and isinstance(prompt, list):166            batch_size = len(prompt)167        else:168            batch_size = prompt_embeds.shape[0]169 170        if prompt_embeds is None:171            text_inputs = self.tokenizer(172                prompt,173                padding="max_length",174                max_length=self.tokenizer.model_max_length,175                truncation=True,176                return_tensors="pt",177            )178            text_input_ids = text_inputs.input_ids179            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids180 181            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(182                text_input_ids, untruncated_ids183            ):184                removed_text = self.tokenizer.batch_decode(185                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]186                )187                logger.warning(188                    "The following part of your input was truncated because CLIP can only handle sequences up to"189                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"190                )191 192            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:193                attention_mask = text_inputs.attention_mask.to(device)194            else:195                attention_mask = None196 197            prompt_embeds = self.text_encoder(198                text_input_ids.to(device),199                attention_mask=attention_mask,200            )201            prompt_embeds = prompt_embeds[0]202 203        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)204 205        bs_embed, seq_len, _ = prompt_embeds.shape206        # duplicate text embeddings for each generation per prompt, using mps friendly method207        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)208        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)209 210        # get unconditional embeddings for classifier free guidance211        if do_classifier_free_guidance and negative_prompt_embeds is None:212            uncond_tokens: List[str]213            if negative_prompt is None:214                uncond_tokens = [""] * batch_size215            elif type(prompt) is not type(negative_prompt):216                raise TypeError(217                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="218                    f" {type(prompt)}."219                )220            elif isinstance(negative_prompt, str):221                uncond_tokens = [negative_prompt]222            elif batch_size != len(negative_prompt):223                raise ValueError(224                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"225                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"226                    " the batch size of `prompt`."227                )228            else:229                uncond_tokens = negative_prompt230 231            max_length = prompt_embeds.shape[1]232            uncond_input = self.tokenizer(233                uncond_tokens,234                padding="max_length",235                max_length=max_length,236                truncation=True,237                return_tensors="pt",238            )239 240            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:241                attention_mask = uncond_input.attention_mask.to(device)242            else:243                attention_mask = None244 245            negative_prompt_embeds = self.text_encoder(246                uncond_input.input_ids.to(device),247                attention_mask=attention_mask,248            )249            negative_prompt_embeds = negative_prompt_embeds[0]250 251        if do_classifier_free_guidance:252            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method253            seq_len = negative_prompt_embeds.shape[1]254 255            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)256 257            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)258            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)259 260            # For classifier free guidance, we need to do two forward passes.261            # Here we concatenate the unconditional and text embeddings into a single batch262            # to avoid doing two forward passes263            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])264 265        return prompt_embeds266 267 268    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents269    def decode_latents(self, latents):270        latents = 1 / 0.18215 * latents271        image = self.vae.decode(latents).sample272        image = (image / 2 + 0.5).clamp(0, 1)273        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16274        image = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()275        return image276 277    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):278        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)279        if isinstance(generator, list) and len(generator) != batch_size:280            raise ValueError(281                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"282                f" size of {batch_size}. Make sure the batch size matches the length of the generators."283            )284 285        if latents is None:286            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)287        else:288            latents = latents.to(device)289        290        # scale the initial noise by the standard deviation required by the scheduler291        latents = latents * self.scheduler.init_noise_sigma292        return latents293 294    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs295    def prepare_extra_step_kwargs(self, generator, eta):296        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature297        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.298        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502299        # and should be between [0, 1]300 301        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())302        extra_step_kwargs = {}303        if accepts_eta:304            extra_step_kwargs["eta"] = eta305 306        # check if the scheduler accepts generator307        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())308        if accepts_generator:309            extra_step_kwargs["generator"] = generator310        return extra_step_kwargs311 312    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker313    def run_safety_checker(self, image, device, dtype):314        if self.safety_checker is not None:315            safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)316            image, has_nsfw_concept = self.safety_checker(317                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)318            )319        else:320            has_nsfw_concept = None321        return image, has_nsfw_concept322 323