tsi-org/tango
0
1# Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/2 3import inspect4from typing import Any, Callable, Dict, List, Optional, Tuple, Union5 6import numpy as np7import PIL.Image8import torch9from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer10 11from diffusers import AutoencoderKL, ControlNetModel, DiffusionPipeline, UNet2DConditionModel, logging12from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker13from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel14from diffusers.schedulers import KarrasDiffusionSchedulers15from diffusers.utils import (16 PIL_INTERPOLATION,17 is_accelerate_available,18 is_accelerate_version,19 randn_tensor,20 replace_example_docstring,21)22 23 24logger = logging.get_logger(__name__) # pylint: disable=invalid-name25 26EXAMPLE_DOC_STRING = """27 Examples:28 ```py29 >>> import numpy as np30 >>> import torch31 >>> from PIL import Image32 >>> from diffusers import ControlNetModel, UniPCMultistepScheduler33 >>> from diffusers.utils import load_image34 35 >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")36 37 >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)38 39 >>> pipe_controlnet = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(40 "runwayml/stable-diffusion-v1-5",41 controlnet=controlnet,42 safety_checker=None,43 torch_dtype=torch.float1644 )45 46 >>> pipe_controlnet.scheduler = UniPCMultistepScheduler.from_config(pipe_controlnet.scheduler.config)47 >>> pipe_controlnet.enable_xformers_memory_efficient_attention()48 >>> pipe_controlnet.enable_model_cpu_offload()49 50 # using image with edges for our canny controlnet51 >>> control_image = load_image(52 "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/vermeer_canny_edged.png")53 54 55 >>> result_img = pipe_controlnet(controlnet_conditioning_image=control_image,56 image=input_image,57 prompt="an android robot, cyberpank, digitl art masterpiece",58 num_inference_steps=20).images[0]59 60 >>> result_img.show()61 ```62"""63 64 65def prepare_image(image):66 if isinstance(image, torch.Tensor):67 # Batch single image68 if image.ndim == 3:69 image = image.unsqueeze(0)70 71 image = image.to(dtype=torch.float32)72 else:73 # preprocess image74 if isinstance(image, (PIL.Image.Image, np.ndarray)):75 image = [image]76 77 if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):78 image = [np.array(i.convert("RGB"))[None, :] for i in image]79 image = np.concatenate(image, axis=0)80 elif isinstance(image, list) and isinstance(image[0], np.ndarray):81 image = np.concatenate([i[None, :] for i in image], axis=0)82 83 image = image.transpose(0, 3, 1, 2)84 image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.085 86 return image87 88 89def prepare_controlnet_conditioning_image(90 controlnet_conditioning_image,91 width,92 height,93 batch_size,94 num_images_per_prompt,95 device,96 dtype,97 do_classifier_free_guidance,98):99 if not isinstance(controlnet_conditioning_image, torch.Tensor):100 if isinstance(controlnet_conditioning_image, PIL.Image.Image):101 controlnet_conditioning_image = [controlnet_conditioning_image]102 103 if isinstance(controlnet_conditioning_image[0], PIL.Image.Image):104 controlnet_conditioning_image = [105 np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"]))[None, :]106 for i in controlnet_conditioning_image107 ]108 controlnet_conditioning_image = np.concatenate(controlnet_conditioning_image, axis=0)109 controlnet_conditioning_image = np.array(controlnet_conditioning_image).astype(np.float32) / 255.0110 controlnet_conditioning_image = controlnet_conditioning_image.transpose(0, 3, 1, 2)111 controlnet_conditioning_image = torch.from_numpy(controlnet_conditioning_image)112 elif isinstance(controlnet_conditioning_image[0], torch.Tensor):113 controlnet_conditioning_image = torch.cat(controlnet_conditioning_image, dim=0)114 115 image_batch_size = controlnet_conditioning_image.shape[0]116 117 if image_batch_size == 1:118 repeat_by = batch_size119 else:120 # image batch size is the same as prompt batch size121 repeat_by = num_images_per_prompt122 123 controlnet_conditioning_image = controlnet_conditioning_image.repeat_interleave(repeat_by, dim=0)124 125 controlnet_conditioning_image = controlnet_conditioning_image.to(device=device, dtype=dtype)126 127 if do_classifier_free_guidance:128 controlnet_conditioning_image = torch.cat([controlnet_conditioning_image] * 2)129 130 return controlnet_conditioning_image131 132 133class StableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline):134 """135 Inspired by: https://github.com/haofanwang/ControlNet-for-Diffusers/136 """137 138 _optional_components = ["safety_checker", "feature_extractor"]139 140 def __init__(141 self,142 vae: AutoencoderKL,143 text_encoder: CLIPTextModel,144 tokenizer: CLIPTokenizer,145 unet: UNet2DConditionModel,146 controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],147 scheduler: KarrasDiffusionSchedulers,148 safety_checker: StableDiffusionSafetyChecker,149 feature_extractor: CLIPImageProcessor,150 requires_safety_checker: bool = True,151 ):152 super().__init__()153 154 if safety_checker is None and requires_safety_checker:155 logger.warning(156 f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"157 " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"158 " results in services or applications open to the public. Both the diffusers team and Hugging Face"159 " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"160 " it only for use-cases that involve analyzing network behavior or auditing its results. For more"161 " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."162 )163 164 if safety_checker is not None and feature_extractor is None:165 raise ValueError(166 "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"167 " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."168 )169 170 if isinstance(controlnet, (list, tuple)):171 controlnet = MultiControlNetModel(controlnet)172 173 self.register_modules(174 vae=vae,175 text_encoder=text_encoder,176 tokenizer=tokenizer,177 unet=unet,178 controlnet=controlnet,179 scheduler=scheduler,180 safety_checker=safety_checker,181 feature_extractor=feature_extractor,182 )183 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)184 self.register_to_config(requires_safety_checker=requires_safety_checker)185 186 def enable_vae_slicing(self):187 r"""188 Enable sliced VAE decoding.189 190 When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several191 steps. This is useful to save some memory and allow larger batch sizes.192 """193 self.vae.enable_slicing()194 195 def disable_vae_slicing(self):196 r"""197 Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to198 computing decoding in one step.199 """200 self.vae.disable_slicing()201 202 def enable_sequential_cpu_offload(self, gpu_id=0):203 r"""204 Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,205 text_encoder, vae, controlnet, and safety checker have their state dicts saved to CPU and then are moved to a206 `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.207 Note that offloading happens on a submodule basis. Memory savings are higher than with208 `enable_model_cpu_offload`, but performance is lower.209 """210 if is_accelerate_available():211 from accelerate import cpu_offload212 else:213 raise ImportError("Please install accelerate via `pip install accelerate`")214 215 device = torch.device(f"cuda:{gpu_id}")216 217 for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.controlnet]:218 cpu_offload(cpu_offloaded_model, device)219 220 if self.safety_checker is not None:221 cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)222 223 def enable_model_cpu_offload(self, gpu_id=0):224 r"""225 Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared226 to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`227 method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with228 `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.229 """230 if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):231 from accelerate import cpu_offload_with_hook232 else:233 raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")234 235 device = torch.device(f"cuda:{gpu_id}")236 237 hook = None238 for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:239 _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)240 241 if self.safety_checker is not None:242 # the safety checker can offload the vae again243 _, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook)244 245 # control net hook has be manually offloaded as it alternates with unet246 cpu_offload_with_hook(self.controlnet, device)247 248 # We'll offload the last model manually.249 self.final_offload_hook = hook250 251 @property252 def _execution_device(self):253 r"""254 Returns the device on which the pipeline's models will be executed. After calling255 `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module256 hooks.257 """258 if not hasattr(self.unet, "_hf_hook"):259 return self.device260 for module in self.unet.modules():261 if (262 hasattr(module, "_hf_hook")263 and hasattr(module._hf_hook, "execution_device")264 and module._hf_hook.execution_device is not None265 ):266 return torch.device(module._hf_hook.execution_device)267 return self.device268 269 def _encode_prompt(270 self,271 prompt,272 device,273 num_images_per_prompt,274 do_classifier_free_guidance,275 negative_prompt=None,276 prompt_embeds: Optional[torch.FloatTensor] = None,277 negative_prompt_embeds: Optional[torch.FloatTensor] = None,278 ):279 r"""280 Encodes the prompt into text encoder hidden states.281 282 Args:283 prompt (`str` or `List[str]`, *optional*):284 prompt to be encoded285 device: (`torch.device`):286 torch device287 num_images_per_prompt (`int`):288 number of images that should be generated per prompt289 do_classifier_free_guidance (`bool`):290 whether to use classifier free guidance or not291 negative_prompt (`str` or `List[str]`, *optional*):292 The prompt or prompts not to guide the image generation. If not defined, one has to pass293 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).294 prompt_embeds (`torch.FloatTensor`, *optional*):295 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not296 provided, text embeddings will be generated from `prompt` input argument.297 negative_prompt_embeds (`torch.FloatTensor`, *optional*):298 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt299 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input300 argument.301 """302 if prompt is not None and isinstance(prompt, str):303 batch_size = 1304 elif prompt is not None and isinstance(prompt, list):305 batch_size = len(prompt)306 else:307 batch_size = prompt_embeds.shape[0]308 309 if prompt_embeds is None:310 text_inputs = self.tokenizer(311 prompt,312 padding="max_length",313 max_length=self.tokenizer.model_max_length,314 truncation=True,315 return_tensors="pt",316 )317 text_input_ids = text_inputs.input_ids318 untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids319 320 if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(321 text_input_ids, untruncated_ids322 ):323 removed_text = self.tokenizer.batch_decode(324 untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]325 )326 logger.warning(327 "The following part of your input was truncated because CLIP can only handle sequences up to"328 f" {self.tokenizer.model_max_length} tokens: {removed_text}"329 )330 331 if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:332 attention_mask = text_inputs.attention_mask.to(device)333 else:334 attention_mask = None335 336 prompt_embeds = self.text_encoder(337 text_input_ids.to(device),338 attention_mask=attention_mask,339 )340 prompt_embeds = prompt_embeds[0]341 342 prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)343 344 bs_embed, seq_len, _ = prompt_embeds.shape345 # duplicate text embeddings for each generation per prompt, using mps friendly method346 prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)347 prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)348 349 # get unconditional embeddings for classifier free guidance350 if do_classifier_free_guidance and negative_prompt_embeds is None:351 uncond_tokens: List[str]352 if negative_prompt is None:353 uncond_tokens = [""] * batch_size354 elif type(prompt) is not type(negative_prompt):355 raise TypeError(356 f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="357 f" {type(prompt)}."358 )359 elif isinstance(negative_prompt, str):360 uncond_tokens = [negative_prompt]361 elif batch_size != len(negative_prompt):362 raise ValueError(363 f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"364 f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"365 " the batch size of `prompt`."366 )367 else:368 uncond_tokens = negative_prompt369 370 max_length = prompt_embeds.shape[1]371 uncond_input = self.tokenizer(372 uncond_tokens,373 padding="max_length",374 max_length=max_length,375 truncation=True,376 return_tensors="pt",377 )378 379 if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:380 attention_mask = uncond_input.attention_mask.to(device)381 else:382 attention_mask = None383 384 negative_prompt_embeds = self.text_encoder(385 uncond_input.input_ids.to(device),386 attention_mask=attention_mask,387 )388 negative_prompt_embeds = negative_prompt_embeds[0]389 390 if do_classifier_free_guidance:391 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method392 seq_len = negative_prompt_embeds.shape[1]393 394 negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)395 396 negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)397 negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)398 399 # For classifier free guidance, we need to do two forward passes.400 # Here we concatenate the unconditional and text embeddings into a single batch401 # to avoid doing two forward passes402 prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])403 404 return prompt_embeds405 406 def run_safety_checker(self, image, device, dtype):407 if self.safety_checker is not None:408 safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)409 image, has_nsfw_concept = self.safety_checker(410 images=image, clip_input=safety_checker_input.pixel_values.to(dtype)411 )412 else:413 has_nsfw_concept = None414 return image, has_nsfw_concept415 416 def decode_latents(self, latents):417 latents = 1 / self.vae.config.scaling_factor * latents418 image = self.vae.decode(latents).sample419 image = (image / 2 + 0.5).clamp(0, 1)420 # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16421 image = image.cpu().permute(0, 2, 3, 1).float().numpy()422 return image423 424 def prepare_extra_step_kwargs(self, generator, eta):425 # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature426 # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.427 # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502428 # and should be between [0, 1]429 430 accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())431 extra_step_kwargs = {}432 if accepts_eta:433 extra_step_kwargs["eta"] = eta434 435 # check if the scheduler accepts generator436 accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())437 if accepts_generator:438 extra_step_kwargs["generator"] = generator439 return extra_step_kwargs440 441 def check_controlnet_conditioning_image(self, image, prompt, prompt_embeds):442 image_is_pil = isinstance(image, PIL.Image.Image)443 image_is_tensor = isinstance(image, torch.Tensor)444 image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)445 image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)446 447 if not image_is_pil and not image_is_tensor and not image_is_pil_list and not image_is_tensor_list:448 raise TypeError(449 "image must be passed and be one of PIL image, torch tensor, list of PIL images, or list of torch tensors"450 )451 452 if image_is_pil:453 image_batch_size = 1454 elif image_is_tensor:455 image_batch_size = image.shape[0]456 elif image_is_pil_list:457 image_batch_size = len(image)458 elif image_is_tensor_list:459 image_batch_size = len(image)460 else:461 raise ValueError("controlnet condition image is not valid")462 463 if prompt is not None and isinstance(prompt, str):464 prompt_batch_size = 1465 elif prompt is not None and isinstance(prompt, list):466 prompt_batch_size = len(prompt)467 elif prompt_embeds is not None:468 prompt_batch_size = prompt_embeds.shape[0]469 else:470 raise ValueError("prompt or prompt_embeds are not valid")471 472 if image_batch_size != 1 and image_batch_size != prompt_batch_size:473 raise ValueError(474 f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"475 )476 477 def check_inputs(478 self,479 prompt,480 image,481 controlnet_conditioning_image,482 height,483 width,484 callback_steps,485 negative_prompt=None,486 prompt_embeds=None,487 negative_prompt_embeds=None,488 strength=None,489 controlnet_guidance_start=None,490 controlnet_guidance_end=None,491 controlnet_conditioning_scale=None,492 ):493 if height % 8 != 0 or width % 8 != 0:494 raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")495 496 if (callback_steps is None) or (497 callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)498 ):499 raise ValueError(500 f"`callback_steps` has to be a positive integer but is {callback_steps} of type"501 f" {type(callback_steps)}."502 )503 504 if prompt is not None and prompt_embeds is not None:505 raise ValueError(506 f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"507 " only forward one of the two."508 )509 elif prompt is None and prompt_embeds is None:510 raise ValueError(511 "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."512 )513 elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):514 raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")515 516 if negative_prompt is not None and negative_prompt_embeds is not None:517 raise ValueError(518 f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"519 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."520 )521 522 if prompt_embeds is not None and negative_prompt_embeds is not None:523 if prompt_embeds.shape != negative_prompt_embeds.shape:524 raise ValueError(525 "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"526 f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"527 f" {negative_prompt_embeds.shape}."528 )529 530 # check controlnet condition image531 532 if isinstance(self.controlnet, ControlNetModel):533 self.check_controlnet_conditioning_image(controlnet_conditioning_image, prompt, prompt_embeds)534 elif isinstance(self.controlnet, MultiControlNetModel):535 if not isinstance(controlnet_conditioning_image, list):536 raise TypeError("For multiple controlnets: `image` must be type `list`")537 538 if len(controlnet_conditioning_image) != len(self.controlnet.nets):539 raise ValueError(540 "For multiple controlnets: `image` must have the same length as the number of controlnets."541 )542 543 for image_ in controlnet_conditioning_image:544 self.check_controlnet_conditioning_image(image_, prompt, prompt_embeds)545 else:546 assert False547 548 # Check `controlnet_conditioning_scale`549 550 if isinstance(self.controlnet, ControlNetModel):551 if not isinstance(controlnet_conditioning_scale, float):552 raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")553 elif isinstance(self.controlnet, MultiControlNetModel):554 if isinstance(controlnet_conditioning_scale, list) and len(controlnet_conditioning_scale) != len(555 self.controlnet.nets556 ):557 raise ValueError(558 "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"559 " the same length as the number of controlnets"560 )561 else:562 assert False563 564 if isinstance(image, torch.Tensor):565 if image.ndim != 3 and image.ndim != 4:566 raise ValueError("`image` must have 3 or 4 dimensions")567 568 if image.ndim == 3:569 image_batch_size = 1570 image_channels, image_height, image_width = image.shape571 elif image.ndim == 4:572 image_batch_size, image_channels, image_height, image_width = image.shape573 else:574 assert False575 576 if image_channels != 3:577 raise ValueError("`image` must have 3 channels")578 579 if image.min() < -1 or image.max() > 1:580 raise ValueError("`image` should be in range [-1, 1]")581 582 if self.vae.config.latent_channels != self.unet.config.in_channels:583 raise ValueError(584 f"The config of `pipeline.unet` expects {self.unet.config.in_channels} but received"585 f" latent channels: {self.vae.config.latent_channels},"586 f" Please verify the config of `pipeline.unet` and the `pipeline.vae`"587 )588 589 if strength < 0 or strength > 1:590 raise ValueError(f"The value of `strength` should in [0.0, 1.0] but is {strength}")591 592 if controlnet_guidance_start < 0 or controlnet_guidance_start > 1:593 raise ValueError(594 f"The value of `controlnet_guidance_start` should in [0.0, 1.0] but is {controlnet_guidance_start}"595 )596 597 if controlnet_guidance_end < 0 or controlnet_guidance_end > 1:598 raise ValueError(599 f"The value of `controlnet_guidance_end` should in [0.0, 1.0] but is {controlnet_guidance_end}"600 )601 602 if controlnet_guidance_start > controlnet_guidance_end:603 raise ValueError(604 "The value of `controlnet_guidance_start` should be less than `controlnet_guidance_end`, but got"605 f" `controlnet_guidance_start` {controlnet_guidance_start} >= `controlnet_guidance_end` {controlnet_guidance_end}"606 )607 608 def get_timesteps(self, num_inference_steps, strength, device):609 # get the original timestep using init_timestep610 init_timestep = min(int(num_inference_steps * strength), num_inference_steps)611 612 t_start = max(num_inference_steps - init_timestep, 0)613 timesteps = self.scheduler.timesteps[t_start:]614 615 return timesteps, num_inference_steps - t_start616 617 def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):618 if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):619 raise ValueError(620 f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"621 )622 623 image = image.to(device=device, dtype=dtype)624 625 batch_size = batch_size * num_images_per_prompt626 if isinstance(generator, list) and len(generator) != batch_size:627 raise ValueError(628 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"629 f" size of {batch_size}. Make sure the batch size matches the length of the generators."630 )631 632 if isinstance(generator, list):633 init_latents = [634 self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)635 ]636 init_latents = torch.cat(init_latents, dim=0)637 else:638 init_latents = self.vae.encode(image).latent_dist.sample(generator)639 640 init_latents = self.vae.config.scaling_factor * init_latents641 642 if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:643 raise ValueError(644 f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."645 )646 else:647 init_latents = torch.cat([init_latents], dim=0)648 649 shape = init_latents.shape650 noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)651 652 # get latents653 init_latents = self.scheduler.add_noise(init_latents, noise, timestep)654 latents = init_latents655 656 return latents657 658 def _default_height_width(self, height, width, image):659 if isinstance(image, list):660 image = image[0]661 662 if height is None:663 if isinstance(image, PIL.Image.Image):664 height = image.height665 elif isinstance(image, torch.Tensor):666 height = image.shape[3]667 668 height = (height // 8) * 8 # round down to nearest multiple of 8669 670 if width is None:671 if isinstance(image, PIL.Image.Image):672 width = image.width673 elif isinstance(image, torch.Tensor):674 width = image.shape[2]675 676 width = (width // 8) * 8 # round down to nearest multiple of 8677 678 return height, width679 680 @torch.no_grad()681 @replace_example_docstring(EXAMPLE_DOC_STRING)682 def __call__(683 self,684 prompt: Union[str, List[str]] = None,685 image: Union[torch.Tensor, PIL.Image.Image] = None,686 controlnet_conditioning_image: Union[687 torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]688 ] = None,689 strength: float = 0.8,690 height: Optional[int] = None,691 width: Optional[int] = None,692 num_inference_steps: int = 50,693 guidance_scale: float = 7.5,694 negative_prompt: Optional[Union[str, List[str]]] = None,695 num_images_per_prompt: Optional[int] = 1,696 eta: float = 0.0,697 generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,698 latents: Optional[torch.FloatTensor] = None,699 prompt_embeds: Optional[torch.FloatTensor] = None,700 negative_prompt_embeds: Optional[torch.FloatTensor] = None,701 output_type: Optional[str] = "pil",702 return_dict: bool = True,703 callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,704 callback_steps: int = 1,705 cross_attention_kwargs: Optional[Dict[str, Any]] = None,706 controlnet_conditioning_scale: Union[float, List[float]] = 1.0,707 controlnet_guidance_start: float = 0.0,708 controlnet_guidance_end: float = 1.0,709 ):710 r"""711 Function invoked when calling the pipeline for generation.712 713 Args:714 prompt (`str` or `List[str]`, *optional*):715 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.716 instead.717 image (`torch.Tensor` or `PIL.Image.Image`):718 `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will719 be masked out with `mask_image` and repainted according to `prompt`.720 controlnet_conditioning_image (`torch.FloatTensor`, `PIL.Image.Image`, `List[torch.FloatTensor]` or `List[PIL.Image.Image]`):721 The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If722 the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. PIL.Image.Image` can723 also be accepted as an image. The control image is automatically resized to fit the output image.724 strength (`float`, *optional*):725 Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`726 will be used as a starting point, adding more noise to it the larger the `strength`. The number of727 denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will728 be maximum and the denoising process will run for the full number of iterations specified in729 `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.730 height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):731 The height in pixels of the generated image.732 width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):733 The width in pixels of the generated image.734 num_inference_steps (`int`, *optional*, defaults to 50):735 The number of denoising steps. More denoising steps usually lead to a higher quality image at the736 expense of slower inference.737 guidance_scale (`float`, *optional*, defaults to 7.5):738 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).739 `guidance_scale` is defined as `w` of equation 2. of [Imagen740 Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >741 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,742 usually at the expense of lower image quality.743 negative_prompt (`str` or `List[str]`, *optional*):744 The prompt or prompts not to guide the image generation. If not defined, one has to pass745 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).746 num_images_per_prompt (`int`, *optional*, defaults to 1):747 The number of images to generate per prompt.748 eta (`float`, *optional*, defaults to 0.0):749 Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to750 [`schedulers.DDIMScheduler`], will be ignored for others.751 generator (`torch.Generator` or `List[torch.Generator]`, *optional*):752 One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)753 to make generation deterministic.754 latents (`torch.FloatTensor`, *optional*):755 Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image756 generation. Can be used to tweak the same generation with different prompts. If not provided, a latents757 tensor will ge generated by sampling using the supplied random `generator`.758 prompt_embeds (`torch.FloatTensor`, *optional*):759 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not760 provided, text embeddings will be generated from `prompt` input argument.761 negative_prompt_embeds (`torch.FloatTensor`, *optional*):762 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt763 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input764 argument.765 output_type (`str`, *optional*, defaults to `"pil"`):766 The output format of the generate image. Choose between767 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.768 return_dict (`bool`, *optional*, defaults to `True`):769 Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a770 plain tuple.771 callback (`Callable`, *optional*):772 A function that will be called every `callback_steps` steps during inference. The function will be773 called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.774 callback_steps (`int`, *optional*, defaults to 1):775 The frequency at which the `callback` function will be called. If not specified, the callback will be776 called at every step.777 cross_attention_kwargs (`dict`, *optional*):778 A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under779 `self.processor` in780 [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).781 controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):782 The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added783 to the residual in the original unet.784 controlnet_guidance_start ('float', *optional*, defaults to 0.0):785 The percentage of total steps the controlnet starts applying. Must be between 0 and 1.786 controlnet_guidance_end ('float', *optional*, defaults to 1.0):787 The percentage of total steps the controlnet ends applying. Must be between 0 and 1. Must be greater788 than `controlnet_guidance_start`.789 790 Examples:791 792 Returns:793 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:794 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.795 When returning a tuple, the first element is a list with the generated images, and the second element is a796 list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"797 (nsfw) content, according to the `safety_checker`.798 """799 # 0. Default height and width to unet800 height, width = self._default_height_width(height, width, controlnet_conditioning_image)801 802 # 1. Check inputs. Raise error if not correct803 self.check_inputs(804 prompt,805 image,806 controlnet_conditioning_image,807 height,808 width,809 callback_steps,810 negative_prompt,811 prompt_embeds,812 negative_prompt_embeds,813 strength,814 controlnet_guidance_start,815 controlnet_guidance_end,816 controlnet_conditioning_scale,817 )818 819 # 2. Define call parameters820 if prompt is not None and isinstance(prompt, str):821 batch_size = 1822 elif prompt is not None and isinstance(prompt, list):823 batch_size = len(prompt)824 else:825 batch_size = prompt_embeds.shape[0]826 827 device = self._execution_device828 # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)829 # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`830 # corresponds to doing no classifier free guidance.831 do_classifier_free_guidance = guidance_scale > 1.0832 833 if isinstance(self.controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):834 controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(self.controlnet.nets)835 836 # 3. Encode input prompt837 prompt_embeds = self._encode_prompt(838 prompt,839 device,840 num_images_per_prompt,841 do_classifier_free_guidance,842 negative_prompt,843 prompt_embeds=prompt_embeds,844 negative_prompt_embeds=negative_prompt_embeds,845 )846 847 # 4. Prepare image, and controlnet_conditioning_image848 image = prepare_image(image)849 850 # condition image(s)851 if isinstance(self.controlnet, ControlNetModel):852 controlnet_conditioning_image = prepare_controlnet_conditioning_image(853 controlnet_conditioning_image=controlnet_conditioning_image,854 width=width,855 height=height,856 batch_size=batch_size * num_images_per_prompt,857 num_images_per_prompt=num_images_per_prompt,858 device=device,859 dtype=self.controlnet.dtype,860 do_classifier_free_guidance=do_classifier_free_guidance,861 )862 elif isinstance(self.controlnet, MultiControlNetModel):863 controlnet_conditioning_images = []864 865 for image_ in controlnet_conditioning_image:866 image_ = prepare_controlnet_conditioning_image(867 controlnet_conditioning_image=image_,868 width=width,869 height=height,870 batch_size=batch_size * num_images_per_prompt,871 num_images_per_prompt=num_images_per_prompt,872 device=device,873 dtype=self.controlnet.dtype,874 do_classifier_free_guidance=do_classifier_free_guidance,875 )876 877 controlnet_conditioning_images.append(image_)878 879 controlnet_conditioning_image = controlnet_conditioning_images880 else:881 assert False882 883 # 5. Prepare timesteps884 self.scheduler.set_timesteps(num_inference_steps, device=device)885 timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)886 latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)887 888 # 6. Prepare latent variables889 latents = self.prepare_latents(890 image,891 latent_timestep,892 batch_size,893 num_images_per_prompt,894 prompt_embeds.dtype,895 device,896 generator,897 )898 899 # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline900 extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)901 902 # 8. Denoising loop903 num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order904 with self.progress_bar(total=num_inference_steps) as progress_bar:905 for i, t in enumerate(timesteps):906 # expand the latents if we are doing classifier free guidance907 latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents908 909 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)910 911 # compute the percentage of total steps we are at912 current_sampling_percent = i / len(timesteps)913 914 if (915 current_sampling_percent < controlnet_guidance_start916 or current_sampling_percent > controlnet_guidance_end917 ):918 # do not apply the controlnet919 down_block_res_samples = None920 mid_block_res_sample = None921 else:922 # apply the controlnet923 down_block_res_samples, mid_block_res_sample = self.controlnet(924 latent_model_input,925 t,926 encoder_hidden_states=prompt_embeds,927 controlnet_cond=controlnet_conditioning_image,928 conditioning_scale=controlnet_conditioning_scale,929 return_dict=False,930 )931 932 # predict the noise residual933 noise_pred = self.unet(934 latent_model_input,935 t,936 encoder_hidden_states=prompt_embeds,937 cross_attention_kwargs=cross_attention_kwargs,938 down_block_additional_residuals=down_block_res_samples,939 mid_block_additional_residual=mid_block_res_sample,940 ).sample941 942 # perform guidance943 if do_classifier_free_guidance:944 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)945 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)946 947 # compute the previous noisy sample x_t -> x_t-1948 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample949 950 # call the callback, if provided951 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):952 progress_bar.update()953 if callback is not None and i % callback_steps == 0:954 callback(i, t, latents)955 956 # If we do sequential model offloading, let's offload unet and controlnet957 # manually for max memory savings958 if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:959 self.unet.to("cpu")960 self.controlnet.to("cpu")961 torch.cuda.empty_cache()962 963 if output_type == "latent":964 image = latents965 has_nsfw_concept = None966 elif output_type == "pil":967 # 8. Post-processing968 image = self.decode_latents(latents)969 970 # 9. Run safety checker971 image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)972 973 # 10. Convert to PIL974 image = self.numpy_to_pil(image)975 else:976 # 8. Post-processing977 image = self.decode_latents(latents)978 979 # 9. Run safety checker980 image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)981 982 # Offload last model to CPU983 if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:984 self.final_offload_hook.offload()985 986 if not return_dict:987 return (image, has_nsfw_concept)988 989 return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)990