Push initial pipeline

rct_diffusion_pipeline.py

@@ -0,0 +1,15 @@
+from diffusers import DiffusionPipeline
+import torch
+
+class RCTDiffusionPipeline(DiffusionPipeline):
+    def __init__(self, unet, scheduler):
+        super().__init__()
+        self.register_modules(unet=unet, scheduler=scheduler)
+    
+    def __call__(self):
+        image = torch.randn((1, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size))
+        timestep = 1
+
+        model_output = self.unet(image, timestep).sample
+        scheduler_output = self.scheduler.step(model_output, timestep, image).prev_sample
+        return scheduler_output

test_pipeline.py

@@ -0,0 +1,8 @@
+from rct_diffusion_pipeline import RCTDiffusionPipeline
+from diffusers import DDPMScheduler, UNet2DModel
+
+scheduler = DDPMScheduler()
+unet = UNet2DModel()
+
+pipeline = RCTDiffusionPipeline(unet=unet, scheduler=scheduler)
+output = pipeline()