Wauplin/diffusers-community-pipelines-mirror
064
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