Upload pipeline.py with huggingface_hub

pipeline.py

@@ -0,0 +1,40 @@
+from diffusers import DiffusionPipeline
+from torch.distributions import Categorical
+import torch
+from tqdm import tqdm
+
+class CRISPRDiffuserPipeline(DiffusionPipeline):
+    def __init__(self, unet, scheduler):
+        super().__init__()
+
+        self.register_modules(unet=unet, scheduler=scheduler)
+        self.stationary_sampler1 = Categorical(probs=unet.stationary_sampler1_probs)
+        self.stationary_sampler2 = Categorical(probs=unet.stationary_sampler2_probs)
+
+    @torch.no_grad()
+    def __call__(self, batch, batch_size=1, record_path=False):
+        x1t = self.stationary_sampler1.sample(torch.Size([batch_size]))
+        x2t = self.stationary_sampler2.sample(torch.Size([batch_size]))
+        t = self.scheduler.step_to_time(torch.tensor([self.scheduler.config.num_train_timesteps]))
+        if record_path:
+            x1ts, x2ts, ts = [x1t], [x2t], [t]
+        for timestep in tqdm(self.scheduler.timesteps):
+            if timestep >= t:
+                continue
+            p_theta_0_logit = self.unet(
+                {
+                    "x1t": x1t.to(self.unet.device),
+                    "x2t": x2t.to(self.unet.device),
+                    "t": t.to(self.unet.device)
+                },
+                batch["condition"].to(self.unet.device).expand(batch_size, -1, -1, -1)
+            )["p_theta_0_logit"].cpu()
+            # the scheduler automatically set t = timestep
+            x1t, x2t, t = self.scheduler.step(p_theta_0_logit, x1t, x2t, t, self.stationary_sampler1, self.stationary_sampler2)
+            if record_path:
+                x1ts.append(x1t)
+                x2ts.append(x2t)
+                ts.append(t)
+        if record_path:
+            return x1ts, x2ts, ts
+        return x1t, x2t