diffusers/community-pipelines-mirror

1from typing import Any, Callable, Dict, List, Optional, Union2 3import torch4from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer5 6from diffusers import (7    AutoencoderKL,8    DDIMScheduler,9    DiffusionPipeline,10    LMSDiscreteScheduler,11    PNDMScheduler,12    StableDiffusionPipeline,13    UNet2DConditionModel,14)15from diffusers.pipelines.pipeline_utils import StableDiffusionMixin16from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput17from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker18 19 20pipe1_model_id = "CompVis/stable-diffusion-v1-1"21pipe2_model_id = "CompVis/stable-diffusion-v1-2"22pipe3_model_id = "CompVis/stable-diffusion-v1-3"23pipe4_model_id = "CompVis/stable-diffusion-v1-4"24 25 26class StableDiffusionComparisonPipeline(DiffusionPipeline, StableDiffusionMixin):27    r"""28    Pipeline for parallel comparison of Stable Diffusion v1-v429    This pipeline inherits from DiffusionPipeline and depends on the use of an Auth Token for30    downloading pre-trained checkpoints from Hugging Face Hub.31    If using Hugging Face Hub, pass the Model ID for Stable Diffusion v1.4 as the previous 3 checkpoints will be loaded32    automatically.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    def __init__(55        self,56        vae: AutoencoderKL,57        text_encoder: CLIPTextModel,58        tokenizer: CLIPTokenizer,59        unet: UNet2DConditionModel,60        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],61        safety_checker: StableDiffusionSafetyChecker,62        feature_extractor: CLIPImageProcessor,63        requires_safety_checker: bool = True,64    ):65        super()._init_()66 67        self.pipe1 = StableDiffusionPipeline.from_pretrained(pipe1_model_id)68        self.pipe2 = StableDiffusionPipeline.from_pretrained(pipe2_model_id)69        self.pipe3 = StableDiffusionPipeline.from_pretrained(pipe3_model_id)70        self.pipe4 = StableDiffusionPipeline(71            vae=vae,72            text_encoder=text_encoder,73            tokenizer=tokenizer,74            unet=unet,75            scheduler=scheduler,76            safety_checker=safety_checker,77            feature_extractor=feature_extractor,78            requires_safety_checker=requires_safety_checker,79        )80 81        self.register_modules(pipeline1=self.pipe1, pipeline2=self.pipe2, pipeline3=self.pipe3, pipeline4=self.pipe4)82 83    @property84    def layers(self) -> Dict[str, Any]:85        return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")}86 87    @torch.no_grad()88    def text2img_sd1_1(89        self,90        prompt: Union[str, List[str]],91        height: int = 512,92        width: int = 512,93        num_inference_steps: int = 50,94        guidance_scale: float = 7.5,95        negative_prompt: Optional[Union[str, List[str]]] = None,96        num_images_per_prompt: Optional[int] = 1,97        eta: float = 0.0,98        generator: Optional[torch.Generator] = None,99        latents: Optional[torch.Tensor] = None,100        output_type: Optional[str] = "pil",101        return_dict: bool = True,102        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,103        callback_steps: int = 1,104        **kwargs,105    ):106        return self.pipe1(107            prompt=prompt,108            height=height,109            width=width,110            num_inference_steps=num_inference_steps,111            guidance_scale=guidance_scale,112            negative_prompt=negative_prompt,113            num_images_per_prompt=num_images_per_prompt,114            eta=eta,115            generator=generator,116            latents=latents,117            output_type=output_type,118            return_dict=return_dict,119            callback=callback,120            callback_steps=callback_steps,121            **kwargs,122        )123 124    @torch.no_grad()125    def text2img_sd1_2(126        self,127        prompt: Union[str, List[str]],128        height: int = 512,129        width: int = 512,130        num_inference_steps: int = 50,131        guidance_scale: float = 7.5,132        negative_prompt: Optional[Union[str, List[str]]] = None,133        num_images_per_prompt: Optional[int] = 1,134        eta: float = 0.0,135        generator: Optional[torch.Generator] = None,136        latents: Optional[torch.Tensor] = None,137        output_type: Optional[str] = "pil",138        return_dict: bool = True,139        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,140        callback_steps: int = 1,141        **kwargs,142    ):143        return self.pipe2(144            prompt=prompt,145            height=height,146            width=width,147            num_inference_steps=num_inference_steps,148            guidance_scale=guidance_scale,149            negative_prompt=negative_prompt,150            num_images_per_prompt=num_images_per_prompt,151            eta=eta,152            generator=generator,153            latents=latents,154            output_type=output_type,155            return_dict=return_dict,156            callback=callback,157            callback_steps=callback_steps,158            **kwargs,159        )160 161    @torch.no_grad()162    def text2img_sd1_3(163        self,164        prompt: Union[str, List[str]],165        height: int = 512,166        width: int = 512,167        num_inference_steps: int = 50,168        guidance_scale: float = 7.5,169        negative_prompt: Optional[Union[str, List[str]]] = None,170        num_images_per_prompt: Optional[int] = 1,171        eta: float = 0.0,172        generator: Optional[torch.Generator] = None,173        latents: Optional[torch.Tensor] = None,174        output_type: Optional[str] = "pil",175        return_dict: bool = True,176        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,177        callback_steps: int = 1,178        **kwargs,179    ):180        return self.pipe3(181            prompt=prompt,182            height=height,183            width=width,184            num_inference_steps=num_inference_steps,185            guidance_scale=guidance_scale,186            negative_prompt=negative_prompt,187            num_images_per_prompt=num_images_per_prompt,188            eta=eta,189            generator=generator,190            latents=latents,191            output_type=output_type,192            return_dict=return_dict,193            callback=callback,194            callback_steps=callback_steps,195            **kwargs,196        )197 198    @torch.no_grad()199    def text2img_sd1_4(200        self,201        prompt: Union[str, List[str]],202        height: int = 512,203        width: int = 512,204        num_inference_steps: int = 50,205        guidance_scale: float = 7.5,206        negative_prompt: Optional[Union[str, List[str]]] = None,207        num_images_per_prompt: Optional[int] = 1,208        eta: float = 0.0,209        generator: Optional[torch.Generator] = None,210        latents: Optional[torch.Tensor] = None,211        output_type: Optional[str] = "pil",212        return_dict: bool = True,213        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,214        callback_steps: int = 1,215        **kwargs,216    ):217        return self.pipe4(218            prompt=prompt,219            height=height,220            width=width,221            num_inference_steps=num_inference_steps,222            guidance_scale=guidance_scale,223            negative_prompt=negative_prompt,224            num_images_per_prompt=num_images_per_prompt,225            eta=eta,226            generator=generator,227            latents=latents,228            output_type=output_type,229            return_dict=return_dict,230            callback=callback,231            callback_steps=callback_steps,232            **kwargs,233        )234 235    @torch.no_grad()236    def _call_(237        self,238        prompt: Union[str, List[str]],239        height: int = 512,240        width: int = 512,241        num_inference_steps: int = 50,242        guidance_scale: float = 7.5,243        negative_prompt: Optional[Union[str, List[str]]] = None,244        num_images_per_prompt: Optional[int] = 1,245        eta: float = 0.0,246        generator: Optional[torch.Generator] = None,247        latents: Optional[torch.Tensor] = None,248        output_type: Optional[str] = "pil",249        return_dict: bool = True,250        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,251        callback_steps: int = 1,252        **kwargs,253    ):254        r"""255        Function invoked when calling the pipeline for generation. This function will generate 4 results as part256        of running all the 4 pipelines for SD1.1-1.4 together in a serial-processing, parallel-invocation fashion.257        Args:258            prompt (`str` or `List[str]`):259                The prompt or prompts to guide the image generation.260            height (`int`, optional, defaults to 512):261                The height in pixels of the generated image.262            width (`int`, optional, defaults to 512):263                The width in pixels of the generated image.264            num_inference_steps (`int`, optional, defaults to 50):265                The number of denoising steps. More denoising steps usually lead to a higher quality image at the266                expense of slower inference.267            guidance_scale (`float`, optional, defaults to 7.5):268                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).269                `guidance_scale` is defined as `w` of equation 2. of [Imagen270                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >271                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,272                usually at the expense of lower image quality.273            eta (`float`, optional, defaults to 0.0):274                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to275                [`schedulers.DDIMScheduler`], will be ignored for others.276            generator (`torch.Generator`, optional):277                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation278                deterministic.279            latents (`torch.Tensor`, optional):280                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image281                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents282                tensor will ge generated by sampling using the supplied random `generator`.283            output_type (`str`, optional, defaults to `"pil"`):284                The output format of the generate image. Choose between285                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.286            return_dict (`bool`, optional, defaults to `True`):287                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a288                plain tuple.289        Returns:290            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:291            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.292            When returning a tuple, the first element is a list with the generated images, and the second element is a293            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"294            (nsfw) content, according to the `safety_checker`.295        """296 297        device = "cuda" if torch.cuda.is_available() else "cpu"298        self.to(device)299 300        # Checks if the height and width are divisible by 8 or not301        if height % 8 != 0 or width % 8 != 0:302            raise ValueError(f"`height` and `width` must be divisible by 8 but are {height} and {width}.")303 304        # Get first result from Stable Diffusion Checkpoint v1.1305        res1 = self.text2img_sd1_1(306            prompt=prompt,307            height=height,308            width=width,309            num_inference_steps=num_inference_steps,310            guidance_scale=guidance_scale,311            negative_prompt=negative_prompt,312            num_images_per_prompt=num_images_per_prompt,313            eta=eta,314            generator=generator,315            latents=latents,316            output_type=output_type,317            return_dict=return_dict,318            callback=callback,319            callback_steps=callback_steps,320            **kwargs,321        )322 323        # Get first result from Stable Diffusion Checkpoint v1.2324        res2 = self.text2img_sd1_2(325            prompt=prompt,326            height=height,327            width=width,328            num_inference_steps=num_inference_steps,329            guidance_scale=guidance_scale,330            negative_prompt=negative_prompt,331            num_images_per_prompt=num_images_per_prompt,332            eta=eta,333            generator=generator,334            latents=latents,335            output_type=output_type,336            return_dict=return_dict,337            callback=callback,338            callback_steps=callback_steps,339            **kwargs,340        )341 342        # Get first result from Stable Diffusion Checkpoint v1.3343        res3 = self.text2img_sd1_3(344            prompt=prompt,345            height=height,346            width=width,347            num_inference_steps=num_inference_steps,348            guidance_scale=guidance_scale,349            negative_prompt=negative_prompt,350            num_images_per_prompt=num_images_per_prompt,351            eta=eta,352            generator=generator,353            latents=latents,354            output_type=output_type,355            return_dict=return_dict,356            callback=callback,357            callback_steps=callback_steps,358            **kwargs,359        )360 361        # Get first result from Stable Diffusion Checkpoint v1.4362        res4 = self.text2img_sd1_4(363            prompt=prompt,364            height=height,365            width=width,366            num_inference_steps=num_inference_steps,367            guidance_scale=guidance_scale,368            negative_prompt=negative_prompt,369            num_images_per_prompt=num_images_per_prompt,370            eta=eta,371            generator=generator,372            latents=latents,373            output_type=output_type,374            return_dict=return_dict,375            callback=callback,376            callback_steps=callback_steps,377            **kwargs,378        )379 380        # Get all result images into a single list and pass it via StableDiffusionPipelineOutput for final result381        return StableDiffusionPipelineOutput([res1[0], res2[0], res3[0], res4[0]])382