import torch
import math
import numpy as np
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, UniPCMultistepScheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from typing import Union, Optional, List, Callable, Dict, Any, Tuple
from momentum_scheduler import (
GHVBScheduler,
PLMSWithHBScheduler,
PLMSWithNTScheduler,
MomentumDPMSolverMultistepScheduler,
MomentumUniPCMultistepScheduler,
)
available_solvers = {
"GHVB": GHVBScheduler,
"PLMS_HB": PLMSWithHBScheduler,
"PLMS_NT": PLMSWithNTScheduler,
"DPM-Solver++": MomentumDPMSolverMultistepScheduler,
"UniPC": MomentumUniPCMultistepScheduler,
}
def get_momentum_number(order, beta):
out = order if beta == 1.0 else order - (1 - beta)
return out
def setup_scheduler(pipe, scheduler, momentum_type="Polyak's heavy ball", order=4.0, beta=1.0, original_config=None):
assert original_config is not None
if scheduler in ["DPM-Solver++", "UniPC"]:
if momentum_type in ["Nesterov"]:
raise NotImplementedError(f"{scheduler} w/ Nesterov is not implemented.")
pipe.scheduler = available_solvers[scheduler].from_config(original_config)
pipe.scheduler.initialize_momentum(beta=beta)
elif scheduler in ["PLMS"]:
momentum_number = get_momentum_number(order, beta)
method = "PLMS_HB" if momentum_type == "Polyak's heavy ball" else "PLMS_NT"
pipe.scheduler = DPMSolverMultistepScheduler.from_config(original_config)
pipe.init_scheduler(method=method, order=momentum_number)
pipe.clear_scheduler()
elif scheduler in ["GHVB"]:
momentum_number = get_momentum_number(order, beta)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(original_config)
pipe.init_scheduler(method="GHVB", order=momentum_number)
pipe.clear_scheduler()
return pipe
class CustomPipeline(StableDiffusionPipeline):
def clear_scheduler(self):
self.scheduler_uncond.clear()
self.scheduler_text.clear()
def init_scheduler(self, method, order):
# equivalent to not applied numerical operator splitting since orders are the same
self.scheduler_uncond = available_solvers[method](self.scheduler, order)
self.scheduler_text = available_solvers[method](self.scheduler, order)
def get_noise(self, latents, prompt_embeds, guidance_scale, t, do_classifier_free_guidance):
# expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
# predict the noise residual
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
grads_a = guidance_scale * (noise_pred_text - noise_pred_uncond)
return noise_pred_uncond, grads_a
def denoising_step(
self,
latents,
prompt_embeds,
guidance_scale,
t,
do_classifier_free_guidance,
method,
extra_step_kwargs,
):
noise_pred_uncond, grads_a = self.get_noise(
latents, prompt_embeds, guidance_scale, t, do_classifier_free_guidance
)
if method in ["dpm", "unipc"]:
latents = self.scheduler.step(noise_pred_uncond + grads_a, t, latents, **extra_step_kwargs).prev_sample
elif method in ["hb", "ghvb", "nt"]:
latents = self.scheduler_uncond.step(noise_pred_uncond, t, latents, output_mode="scale")
latents = self.scheduler_text.step(grads_a, t, latents, output_mode='back')
else:
raise NotImplementedError
return latents
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
method="ghvb",
):
# 0. Default height and width to unet
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
)
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale > 1.0
# 3. Encode input prompt
prompt_embeds = self._encode_prompt(
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
)
# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
# print(timesteps)
# 5. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 7. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
latents = self.denoising_step(
latents,
prompt_embeds,
guidance_scale,
t,
do_classifier_free_guidance,
method,
extra_step_kwargs,
)
# call the callback, if provided
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
callback(i, t, latents)
if output_type == "latent":
image = latents
has_nsfw_concept = None
elif output_type == "pil":
# 8. Post-processing
image = self.decode_latents(latents)
# 9. Run safety checker
# image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
has_nsfw_concept = False
# 10. Convert to PIL
image = self.numpy_to_pil(image)
else:
# 8. Post-processing
image = self.decode_latents(latents)
# 9. Run safety checker
# image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
has_nsfw_concept = False
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
def generate(self, params):
params["output_type"] = "latent"
ori_latents = self.__call__(**params)["images"]
with torch.no_grad():
latents = torch.clone(ori_latents)
image = self.decode_latents(latents)
image = self.numpy_to_pil(image)[0]
return image, ori_latents