Upload pipeline.py

pipeline.py

@@ -0,0 +1,83 @@
+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)
+
+        # TODO: Implement greedy search on FID
+
+    def search_previous_time(
+            self, time, eps: float = 0.002, T: float = 80.0
+    ):
+        return (2 * time + eps) / 3
+