v0.3.0

pipeline.py

@@ -0,0 +1,76 @@
+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) * 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=sample.device, generator=generator
+                )
+
+            out = model(
+                sample, torch.tensor([time], device=sample.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)
+
+    # TODO: Implement greedy search on FID
+    def search_previous_time(
+        self, time, time_min: float = 0.002, time_max: float = 80.0
+    ):
+        return (2 * time + time_min) / 3