Spaces:
Runtime error
Runtime error
| import inspect | |
| import warnings | |
| from typing import Callable, List, Optional, Tuple, Union | |
| import numpy as np | |
| import torch | |
| from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer | |
| from diffusers.models import AutoencoderKL, UNet2DConditionModel | |
| from diffusers.pipeline_utils import DiffusionPipeline | |
| from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler | |
| from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput | |
| from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker | |
| class StableDiffusionPipelineWithCallback(DiffusionPipeline): | |
| r""" | |
| Pipeline for text-to-image generation using Stable Diffusion. | |
| ** based on https://github.com/huggingface/diffusers/pull/521/files#diff-ab952f41078da66b9fcbbd913b419f8c334badceefac03a5f7edcd6dd986a8ef ** | |
| This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the | |
| library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) | |
| Args: | |
| vae ([`AutoencoderKL`]): | |
| Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. | |
| text_encoder ([`CLIPTextModel`]): | |
| Frozen text-encoder. Stable Diffusion uses the text portion of | |
| [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically | |
| the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. | |
| tokenizer (`CLIPTokenizer`): | |
| Tokenizer of class | |
| [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). | |
| unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. | |
| scheduler ([`SchedulerMixin`]): | |
| A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of | |
| [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. | |
| safety_checker ([`StableDiffusionSafetyChecker`]): | |
| Classification module that estimates whether generated images could be considered offsensive or harmful. | |
| Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details. | |
| feature_extractor ([`CLIPFeatureExtractor`]): | |
| Model that extracts features from generated images to be used as inputs for the `safety_checker`. | |
| """ | |
| def __init__( | |
| self, | |
| vae: AutoencoderKL, | |
| text_encoder: CLIPTextModel, | |
| tokenizer: CLIPTokenizer, | |
| unet: UNet2DConditionModel, | |
| scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], | |
| safety_checker: StableDiffusionSafetyChecker, | |
| feature_extractor: CLIPFeatureExtractor, | |
| ): | |
| super().__init__() | |
| scheduler = scheduler.set_format("pt") | |
| self.register_modules( | |
| vae=vae, | |
| text_encoder=text_encoder, | |
| tokenizer=tokenizer, | |
| unet=unet, | |
| scheduler=scheduler, | |
| safety_checker=safety_checker, | |
| feature_extractor=feature_extractor, | |
| ) | |
| def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): | |
| r""" | |
| Enable sliced attention computation. | |
| When this option is enabled, the attention module will split the input tensor in slices, to compute attention | |
| in several steps. This is useful to save some memory in exchange for a small speed decrease. | |
| Args: | |
| slice_size (`str` or `int`, *optional*, defaults to `"auto"`): | |
| When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If | |
| a number is provided, uses as many slices as `attention_head_dim // slice_size`. In this case, | |
| `attention_head_dim` must be a multiple of `slice_size`. | |
| """ | |
| if slice_size == "auto": | |
| # half the attention head size is usually a good trade-off between | |
| # speed and memory | |
| slice_size = self.unet.config.attention_head_dim // 2 | |
| self.unet.set_attention_slice(slice_size) | |
| def disable_attention_slicing(self): | |
| r""" | |
| Disable sliced attention computation. If `enable_attention_slicing` was previously invoked, this method will go | |
| back to computing attention in one step. | |
| """ | |
| # set slice_size = `None` to disable `attention slicing` | |
| self.enable_attention_slicing(None) | |
| def decode_latents(self, latents: torch.FloatTensor) -> np.ndarray: | |
| r""" | |
| Scale and decode the latent representations into images using the VAE. | |
| Args: | |
| latents (`torch.FloatTensor`): | |
| Latent representations to decode into images. | |
| Returns: | |
| `np.ndarray`: Decoded images. | |
| """ | |
| latents = 1 / 0.18215 * latents | |
| image = self.vae.decode(latents).sample | |
| image = (image / 2 + 0.5).clamp(0, 1) | |
| image = image.cpu().permute(0, 2, 3, 1).numpy() | |
| return image | |
| def run_safety_checker(self, image: np.ndarray) -> Tuple[np.ndarray, List[bool]]: | |
| r""" | |
| Run the safety checker on the generated images. If potential NSFW content was detected, a warning will be | |
| raised and a black image will be returned instead. | |
| Args: | |
| image (`np.ndarray`): | |
| Images to run the safety checker on. | |
| Returns: | |
| `Tuple[np.ndarray, List[bool]]`: The first element contains the images that has been processed by the | |
| safety checker. The second element is a boolean array indicating whether the images contain NSFW content. | |
| """ | |
| safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(self.device) | |
| image, has_nsfw_concept = self.safety_checker(images=image, clip_input=safety_checker_input.pixel_values) | |
| return image, has_nsfw_concept | |
| def __call__( | |
| self, | |
| prompt: Union[str, List[str]], | |
| height: Optional[int] = 512, | |
| width: Optional[int] = 512, | |
| num_inference_steps: Optional[int] = 50, | |
| guidance_scale: Optional[float] = 7.5, | |
| eta: Optional[float] = 0.0, | |
| generator: Optional[torch.Generator] = None, | |
| latents: Optional[torch.FloatTensor] = None, | |
| output_type: Optional[str] = "pil", | |
| return_dict: bool = True, | |
| callback: Optional[ | |
| Callable[[int, np.ndarray, torch.FloatTensor], None] | |
| ] = None, | |
| callback_frequency: Optional[int] = 1, | |
| **kwargs, | |
| ): | |
| r""" | |
| Function invoked when calling the pipeline for generation. | |
| Args: | |
| prompt (`str` or `List[str]`): | |
| The prompt or prompts to guide the image generation. | |
| height (`int`, *optional*, defaults to 512): | |
| The height in pixels of the generated image. | |
| width (`int`, *optional*, defaults to 512): | |
| The width in pixels of the generated image. | |
| num_inference_steps (`int`, *optional*, defaults to 50): | |
| The number of denoising steps. More denoising steps usually lead to a higher quality image at the | |
| expense of slower inference. | |
| guidance_scale (`float`, *optional*, defaults to 7.5): | |
| Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). | |
| `guidance_scale` is defined as `w` of equation 2. of [Imagen | |
| Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > | |
| 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, | |
| usually at the expense of lower image quality. | |
| eta (`float`, *optional*, defaults to 0.0): | |
| Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to | |
| [`schedulers.DDIMScheduler`], will be ignored for others. | |
| generator (`torch.Generator`, *optional*): | |
| A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation | |
| deterministic. | |
| latents (`torch.FloatTensor`, *optional*): | |
| Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image | |
| generation. Can be used to tweak the same generation with different prompts. If not provided, a latents | |
| tensor will ge generated by sampling using the supplied random `generator`. | |
| output_type (`str`, *optional*, defaults to `"pil"`): | |
| The output format of the generate image. Choose between | |
| [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `nd.array`. | |
| return_dict (`bool`, *optional*, defaults to `True`): | |
| Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a | |
| plain tuple. | |
| callback (`Callable`, *optional*): | |
| A function that will be called every `callback_frequency` steps during inference. The function will be | |
| called with the following arguments: `callback(step: int, timestep: np.ndarray, latents: | |
| torch.FloatTensor, image: Union[List[PIL.Image.Image], np.ndarray])`. | |
| callback_frequency (`int`, *optional*, defaults to 1): | |
| The frequency at which the `callback` function will be called. If not specified, the callback will be | |
| called at every step. | |
| Returns: | |
| [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: | |
| [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. | |
| When returning a tuple, the first element is a list with the generated images, and the second element is a | |
| list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" | |
| (nsfw) content, according to the `safety_checker`. | |
| """ | |
| if "torch_device" in kwargs: | |
| device = kwargs.pop("torch_device") | |
| warnings.warn( | |
| "`torch_device` is deprecated as an input argument to `__call__` and will be removed in v0.3.0." | |
| " Consider using `pipe.to(torch_device)` instead." | |
| ) | |
| # Set device as before (to be removed in 0.3.0) | |
| if device is None: | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.to(device) | |
| if isinstance(prompt, str): | |
| batch_size = 1 | |
| elif isinstance(prompt, list): | |
| batch_size = len(prompt) | |
| else: | |
| raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") | |
| if height % 8 != 0 or width % 8 != 0: | |
| raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") | |
| if (callback_frequency is None) or ( | |
| callback_frequency is not None and (not isinstance(callback_frequency, int) or callback_frequency <= 0) | |
| ): | |
| raise ValueError( | |
| f"`callback_frequency` has to be a positive integer but is {callback_frequency} of type" | |
| f" {type(callback_frequency)}." | |
| ) | |
| # get prompt text embeddings | |
| text_input = self.tokenizer( | |
| prompt, | |
| padding="max_length", | |
| max_length=self.tokenizer.model_max_length, | |
| truncation=True, | |
| return_tensors="pt", | |
| ) | |
| text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0] | |
| # 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. | |
| do_classifier_free_guidance = guidance_scale > 1.0 | |
| # get unconditional embeddings for classifier free guidance | |
| if do_classifier_free_guidance: | |
| max_length = text_input.input_ids.shape[-1] | |
| uncond_input = self.tokenizer( | |
| [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt" | |
| ) | |
| uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0] | |
| # For classifier free guidance, we need to do two forward passes. | |
| # Here we concatenate the unconditional and text embeddings into a single batch | |
| # to avoid doing two forward passes | |
| text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) | |
| # get the initial random noise unless the user supplied it | |
| # Unlike in other pipelines, latents need to be generated in the target device | |
| # for 1-to-1 results reproducibility with the CompVis implementation. | |
| # However this currently doesn't work in `mps`. | |
| latents_device = "cpu" if self.device.type == "mps" else self.device | |
| latents_shape = (batch_size, self.unet.in_channels, height // 8, width // 8) | |
| if latents is None: | |
| latents = torch.randn( | |
| latents_shape, | |
| generator=generator, | |
| device=latents_device, | |
| ) | |
| else: | |
| if latents.shape != latents_shape: | |
| raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") | |
| latents = latents.to(self.device) | |
| # set timesteps | |
| accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys()) | |
| extra_set_kwargs = {} | |
| if accepts_offset: | |
| extra_set_kwargs["offset"] = 1 | |
| self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs) | |
| # if we use LMSDiscreteScheduler, let's make sure latents are mulitplied by sigmas | |
| if isinstance(self.scheduler, LMSDiscreteScheduler): | |
| latents = latents * self.scheduler.sigmas[0] | |
| # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature | |
| # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. | |
| # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 | |
| # and should be between [0, 1] | |
| accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) | |
| extra_step_kwargs = {} | |
| if accepts_eta: | |
| extra_step_kwargs["eta"] = eta | |
| for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): | |
| # expand the latents if we are doing classifier free guidance | |
| latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents | |
| if isinstance(self.scheduler, LMSDiscreteScheduler): | |
| sigma = self.scheduler.sigmas[i] | |
| # the model input needs to be scaled to match the continuous ODE formulation in K-LMS | |
| latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) | |
| # predict the noise residual | |
| noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample | |
| # perform guidance | |
| if 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 | |
| if isinstance(self.scheduler, LMSDiscreteScheduler): | |
| latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs).prev_sample | |
| else: | |
| latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample | |
| # call the callback, if provided | |
| if callback is not None and i % callback_frequency == 0: | |
| callback(i, t, latents) | |
| image = self.decode_latents(latents) | |
| image, has_nsfw_concept = self.run_safety_checker(image) | |
| if output_type == "pil": | |
| image = self.numpy_to_pil(image) | |
| if not return_dict: | |
| return (image, has_nsfw_concept) | |
| return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) | |