File size: 2,124 Bytes
8784a1d 8dc6ca5 8784a1d 8dc6ca5 8784a1d 8dc6ca5 8784a1d 8dc6ca5 8784a1d 8dc6ca5 8784a1d 8dc6ca5 8784a1d 8dc6ca5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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,
steps: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
time_min: float = 0.002,
time_max: float = 80.0,
data_std: float = 0.5,
output_type: Optional[str] = "pil",
return_dict: bool = True,
**kwargs,
) -> Union[Tuple, ImagePipelineOutput]:
img_size = self.unet.config.sample_size
shape = (1, 3, img_size, img_size)
model = self.unet
time: float = time_max
sample = randn_tensor(shape, generator=generator, device=self.device) * time
for step in self.progress_bar(range(steps)):
if step > 0:
time = self.search_previous_time(time)
sigma = math.sqrt(time**2 - time_min**2 + 1e-6)
sample = sample + sigma * randn_tensor(sample.shape, device=self.device, generator=generator)
out = model(sample, torch.tensor([time], device=self.device)).sample
skip_coef = data_std**2 / ((time - time_min) ** 2 + data_std**2)
out_coef = data_std * time / (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, time_min: float = 0.002, time_max: float = 80.0):
return (2 * time + time_min) / 3
def cuda(self):
self.to("cuda")
|