|
import math |
|
from typing import List, Optional, Tuple, Union |
|
|
|
import torch |
|
from diffusers import DiffusionPipeline, ImagePipelineOutput, UNet2DModel |
|
from diffusers.utils import randn_tensor |
|
|
|
|
|
class ConsistencyPipeline(DiffusionPipeline): |
|
unet: UNet2DModel |
|
|
|
def __init__( |
|
self, |
|
unet: UNet2DModel, |
|
) -> None: |
|
super().__init__() |
|
self.register_modules(unet=unet) |
|
|
|
@torch.no_grad() |
|
def __call__( |
|
self, |
|
batch_size : int = 1, |
|
num_class: Optional[int] = None, |
|
label_index: Optional[int] = None, |
|
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, |
|
eps: float = 0.002, |
|
T: float = 80.0, |
|
data_std: float = 0.5, |
|
num_inference_steps: int = 1, |
|
output_type: Optional[str] = "pil", |
|
return_dict: bool = True, |
|
**kwargs, |
|
) -> Union[Tuple, ImagePipelineOutput]: |
|
model = self.unet |
|
device = model.device |
|
image_labels = None |
|
if label_index is not None: |
|
assert label_index + 1 <= num_class, 'label_index must <= num_class!' |
|
image_labels = torch.LongTensor([label_index]).repeat(batch_size).to(device) |
|
else: |
|
if num_class is not None: |
|
image_labels = torch.randint(low=0, high=num_class, size=[1]) |
|
image_labels = image_labels.repeat(batch_size).to(device) |
|
img_size = self.unet.config.sample_size |
|
shape = (batch_size, 3, img_size, img_size) |
|
|
|
time: float = T |
|
|
|
sample = randn_tensor(shape, generator=generator, device=device) * time |
|
|
|
for step in self.progress_bar(range(num_inference_steps)): |
|
if step > 0: |
|
time = self.search_previous_time(time) |
|
sigma = math.sqrt(time ** 2 - eps ** 2 + 1e-6) |
|
sample = sample + sigma * randn_tensor( |
|
sample.shape, device=sample.device, generator=generator |
|
) |
|
|
|
out = model(sample, torch.tensor([time], device=sample.device), image_labels).sample |
|
|
|
skip_coef = data_std ** 2 / ((time - eps) ** 2 + data_std ** 2) |
|
out_coef = data_std * (time - eps) / (time ** 2 + data_std ** 2) ** (0.5) |
|
|
|
sample = (sample * skip_coef + out * out_coef).clamp(-1.0, 1.0) |
|
|
|
sample = (sample / 2 + 0.5).clamp(0, 1) |
|
image = sample.cpu().permute(0, 2, 3, 1).numpy() |
|
|
|
if output_type == "pil": |
|
image = self.numpy_to_pil(image) |
|
|
|
if not return_dict: |
|
return (image,) |
|
|
|
return ImagePipelineOutput(images=image) |
|
|
|
|
|
|
|
def search_previous_time( |
|
self, time, eps: float = 0.002, T: float = 80.0 |
|
): |
|
return (2 * time + eps) / 3 |
|
|
|
|
