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.pipeline_utils import StableDiffusionMixin20from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker21from diffusers.utils import deprecate, logging22 23 24logger = logging.get_logger(__name__)  # pylint: disable=invalid-name25 26 27class StableDiffusionMegaPipeline(DiffusionPipeline, StableDiffusionMixin):28    r"""29    Pipeline for text-to-image generation using Stable Diffusion.30 31    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the32    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)33 34    Args:35        vae ([`AutoencoderKL`]):36            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.37        text_encoder ([`CLIPTextModel`]):38            Frozen text-encoder. Stable Diffusion uses the text portion of39            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically40            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.41        tokenizer (`CLIPTokenizer`):42            Tokenizer of class43            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).44        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.45        scheduler ([`SchedulerMixin`]):46            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of47            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].48        safety_checker ([`StableDiffusionMegaSafetyChecker`]):49            Classification module that estimates whether generated images could be considered offensive or harmful.50            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.51        feature_extractor ([`CLIPImageProcessor`]):52            Model that extracts features from generated images to be used as inputs for the `safety_checker`.53    """54 55    _optional_components = ["safety_checker", "feature_extractor"]56 57    def __init__(58        self,59        vae: AutoencoderKL,60        text_encoder: CLIPTextModel,61        tokenizer: CLIPTokenizer,62        unet: UNet2DConditionModel,63        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],64        safety_checker: StableDiffusionSafetyChecker,65        feature_extractor: CLIPImageProcessor,66        requires_safety_checker: bool = True,67    ):68        super().__init__()69        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:70            deprecation_message = (71                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"72                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "73                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"74                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"75                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"76                " file"77            )78            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)79            new_config = dict(scheduler.config)80            new_config["steps_offset"] = 181            scheduler._internal_dict = FrozenDict(new_config)82 83        self.register_modules(84            vae=vae,85            text_encoder=text_encoder,86            tokenizer=tokenizer,87            unet=unet,88            scheduler=scheduler,89            safety_checker=safety_checker,90            feature_extractor=feature_extractor,91        )92        self.register_to_config(requires_safety_checker=requires_safety_checker)93 94    @property95    def components(self) -> Dict[str, Any]:96        return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")}97 98    @torch.no_grad()99    def inpaint(100        self,101        prompt: Union[str, List[str]],102        image: Union[torch.Tensor, PIL.Image.Image],103        mask_image: Union[torch.Tensor, PIL.Image.Image],104        strength: float = 0.8,105        num_inference_steps: Optional[int] = 50,106        guidance_scale: Optional[float] = 7.5,107        negative_prompt: Optional[Union[str, List[str]]] = None,108        num_images_per_prompt: Optional[int] = 1,109        eta: Optional[float] = 0.0,110        generator: Optional[torch.Generator] = None,111        output_type: Optional[str] = "pil",112        return_dict: bool = True,113        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,114        callback_steps: int = 1,115    ):116        # For more information on how this function works, please see: https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion#diffusers.StableDiffusionImg2ImgPipeline117        return StableDiffusionInpaintPipelineLegacy(**self.components)(118            prompt=prompt,119            image=image,120            mask_image=mask_image,121            strength=strength,122            num_inference_steps=num_inference_steps,123            guidance_scale=guidance_scale,124            negative_prompt=negative_prompt,125            num_images_per_prompt=num_images_per_prompt,126            eta=eta,127            generator=generator,128            output_type=output_type,129            return_dict=return_dict,130            callback=callback,131        )132 133    @torch.no_grad()134    def img2img(135        self,136        prompt: Union[str, List[str]],137        image: Union[torch.Tensor, PIL.Image.Image],138        strength: float = 0.8,139        num_inference_steps: Optional[int] = 50,140        guidance_scale: Optional[float] = 7.5,141        negative_prompt: Optional[Union[str, List[str]]] = None,142        num_images_per_prompt: Optional[int] = 1,143        eta: Optional[float] = 0.0,144        generator: Optional[torch.Generator] = None,145        output_type: Optional[str] = "pil",146        return_dict: bool = True,147        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,148        callback_steps: int = 1,149        **kwargs,150    ):151        # For more information on how this function works, please see: https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion#diffusers.StableDiffusionImg2ImgPipeline152        return StableDiffusionImg2ImgPipeline(**self.components)(153            prompt=prompt,154            image=image,155            strength=strength,156            num_inference_steps=num_inference_steps,157            guidance_scale=guidance_scale,158            negative_prompt=negative_prompt,159            num_images_per_prompt=num_images_per_prompt,160            eta=eta,161            generator=generator,162            output_type=output_type,163            return_dict=return_dict,164            callback=callback,165            callback_steps=callback_steps,166        )167 168    @torch.no_grad()169    def text2img(170        self,171        prompt: Union[str, List[str]],172        height: int = 512,173        width: int = 512,174        num_inference_steps: int = 50,175        guidance_scale: float = 7.5,176        negative_prompt: Optional[Union[str, List[str]]] = None,177        num_images_per_prompt: Optional[int] = 1,178        eta: float = 0.0,179        generator: Optional[torch.Generator] = None,180        latents: Optional[torch.Tensor] = None,181        output_type: Optional[str] = "pil",182        return_dict: bool = True,183        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,184        callback_steps: int = 1,185    ):186        # For more information on how this function https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion#diffusers.StableDiffusionPipeline187        return StableDiffusionPipeline(**self.components)(188            prompt=prompt,189            height=height,190            width=width,191            num_inference_steps=num_inference_steps,192            guidance_scale=guidance_scale,193            negative_prompt=negative_prompt,194            num_images_per_prompt=num_images_per_prompt,195            eta=eta,196            generator=generator,197            latents=latents,198            output_type=output_type,199            return_dict=return_dict,200            callback=callback,201            callback_steps=callback_steps,202        )203