diffusers/community-pipelines-mirror

1from typing import Any, Callable, Dict, List, Optional, Union2 3import PIL.Image4import torch5from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer6 7from diffusers import (8    AutoencoderKL,9    DDIMScheduler,10    DiffusionPipeline,11    LMSDiscreteScheduler,12    PNDMScheduler,13    StableDiffusionImg2ImgPipeline,14    StableDiffusionInpaintPipelineLegacy,15    StableDiffusionPipeline,16    UNet2DConditionModel,17)18from diffusers.configuration_utils import FrozenDict19from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker20from diffusers.utils import deprecate, logging21 22 23logger = logging.get_logger(__name__)  # pylint: disable=invalid-name24 25 26class StableDiffusionMegaPipeline(DiffusionPipeline):27    r"""28    Pipeline for text-to-image generation using Stable Diffusion.29 30    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the31    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)32 33    Args:34        vae ([`AutoencoderKL`]):35            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.36        text_encoder ([`CLIPTextModel`]):37            Frozen text-encoder. Stable Diffusion uses the text portion of38            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically39            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.40        tokenizer (`CLIPTokenizer`):41            Tokenizer of class42            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).43        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.44        scheduler ([`SchedulerMixin`]):45            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of46            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].47        safety_checker ([`StableDiffusionMegaSafetyChecker`]):48            Classification module that estimates whether generated images could be considered offensive or harmful.49            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.50        feature_extractor ([`CLIPImageProcessor`]):51            Model that extracts features from generated images to be used as inputs for the `safety_checker`.52    """53 54    _optional_components = ["safety_checker", "feature_extractor"]55 56    def __init__(57        self,58        vae: AutoencoderKL,59        text_encoder: CLIPTextModel,60        tokenizer: CLIPTokenizer,61        unet: UNet2DConditionModel,62        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],63        safety_checker: StableDiffusionSafetyChecker,64        feature_extractor: CLIPImageProcessor,65        requires_safety_checker: bool = True,66    ):67        super().__init__()68        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:69            deprecation_message = (70                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"71                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "72                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"73                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"74                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"75                " file"76            )77            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)78            new_config = dict(scheduler.config)79            new_config["steps_offset"] = 180            scheduler._internal_dict = FrozenDict(new_config)81 82        self.register_modules(83            vae=vae,84            text_encoder=text_encoder,85            tokenizer=tokenizer,86            unet=unet,87            scheduler=scheduler,88            safety_checker=safety_checker,89            feature_extractor=feature_extractor,90        )91        self.register_to_config(requires_safety_checker=requires_safety_checker)92 93    @property94    def components(self) -> Dict[str, Any]:95        return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")}96 97    def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):98        r"""99        Enable sliced attention computation.100 101        When this option is enabled, the attention module will split the input tensor in slices, to compute attention102        in several steps. This is useful to save some memory in exchange for a small speed decrease.103 104        Args:105            slice_size (`str` or `int`, *optional*, defaults to `"auto"`):106                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If107                a number is provided, uses as many slices as `attention_head_dim // slice_size`. In this case,108                `attention_head_dim` must be a multiple of `slice_size`.109        """110        if slice_size == "auto":111            # half the attention head size is usually a good trade-off between112            # speed and memory113            slice_size = self.unet.config.attention_head_dim // 2114        self.unet.set_attention_slice(slice_size)115 116    def disable_attention_slicing(self):117        r"""118        Disable sliced attention computation. If `enable_attention_slicing` was previously invoked, this method will go119        back to computing attention in one step.120        """121        # set slice_size = `None` to disable `attention slicing`122        self.enable_attention_slicing(None)123 124    @torch.no_grad()125    def inpaint(126        self,127        prompt: Union[str, List[str]],128        image: Union[torch.FloatTensor, PIL.Image.Image],129        mask_image: Union[torch.FloatTensor, PIL.Image.Image],130        strength: float = 0.8,131        num_inference_steps: Optional[int] = 50,132        guidance_scale: Optional[float] = 7.5,133        negative_prompt: Optional[Union[str, List[str]]] = None,134        num_images_per_prompt: Optional[int] = 1,135        eta: Optional[float] = 0.0,136        generator: Optional[torch.Generator] = None,137        output_type: Optional[str] = "pil",138        return_dict: bool = True,139        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,140        callback_steps: int = 1,141    ):142        # For more information on how this function works, please see: https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion#diffusers.StableDiffusionImg2ImgPipeline143        return StableDiffusionInpaintPipelineLegacy(**self.components)(144            prompt=prompt,145            image=image,146            mask_image=mask_image,147            strength=strength,148            num_inference_steps=num_inference_steps,149            guidance_scale=guidance_scale,150            negative_prompt=negative_prompt,151            num_images_per_prompt=num_images_per_prompt,152            eta=eta,153            generator=generator,154            output_type=output_type,155            return_dict=return_dict,156            callback=callback,157        )158 159    @torch.no_grad()160    def img2img(161        self,162        prompt: Union[str, List[str]],163        image: Union[torch.FloatTensor, PIL.Image.Image],164        strength: float = 0.8,165        num_inference_steps: Optional[int] = 50,166        guidance_scale: Optional[float] = 7.5,167        negative_prompt: Optional[Union[str, List[str]]] = None,168        num_images_per_prompt: Optional[int] = 1,169        eta: Optional[float] = 0.0,170        generator: Optional[torch.Generator] = None,171        output_type: Optional[str] = "pil",172        return_dict: bool = True,173        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,174        callback_steps: int = 1,175        **kwargs,176    ):177        # For more information on how this function works, please see: https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion#diffusers.StableDiffusionImg2ImgPipeline178        return StableDiffusionImg2ImgPipeline(**self.components)(179            prompt=prompt,180            image=image,181            strength=strength,182            num_inference_steps=num_inference_steps,183            guidance_scale=guidance_scale,184            negative_prompt=negative_prompt,185            num_images_per_prompt=num_images_per_prompt,186            eta=eta,187            generator=generator,188            output_type=output_type,189            return_dict=return_dict,190            callback=callback,191            callback_steps=callback_steps,192        )193 194    @torch.no_grad()195    def text2img(196        self,197        prompt: Union[str, List[str]],198        height: int = 512,199        width: int = 512,200        num_inference_steps: int = 50,201        guidance_scale: float = 7.5,202        negative_prompt: Optional[Union[str, List[str]]] = None,203        num_images_per_prompt: Optional[int] = 1,204        eta: float = 0.0,205        generator: Optional[torch.Generator] = None,206        latents: Optional[torch.FloatTensor] = None,207        output_type: Optional[str] = "pil",208        return_dict: bool = True,209        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,210        callback_steps: int = 1,211    ):212        # For more information on how this function https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion#diffusers.StableDiffusionPipeline213        return StableDiffusionPipeline(**self.components)(214            prompt=prompt,215            height=height,216            width=width,217            num_inference_steps=num_inference_steps,218            guidance_scale=guidance_scale,219            negative_prompt=negative_prompt,220            num_images_per_prompt=num_images_per_prompt,221            eta=eta,222            generator=generator,223            latents=latents,224            output_type=output_type,225            return_dict=return_dict,226            callback=callback,227            callback_steps=callback_steps,228        )229