Create pipeline_spad.py

this is just an experiment, not the perfect spad, more work to do (fix upcoming bugs)

pipeline_spad.py

@@ -0,0 +1,68 @@
+import torch
+from diffusers import AutoencoderKL, DiffusionPipeline
+from transformers import CLIPTextModel, CLIPTokenizer
+from mv_unet import SPADUnetModel
+from diffusers.schedulers import DPMSolverMultistepScheduler
+
+class SPADPipeline(DiffusionPipeline):
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        unet: SPADUnetModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModel,
+        scheduler: DPMSolverMultistepScheduler,
+    ):
+        super().__init__()
+
+        self.vae = vae
+        self.unet = unet
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.scheduler = scheduler
+
+        # make sure all our models are on the same device
+        self.vae.to(self.device)
+        self.unet.to(self.device)
+        self.text_encoder.to(self.device)
+
+    def encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None):
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt"
+        )
+        text_embeddings = self.text_encoder(text_input.input_ids.to(device))[0]
+
+        # we duplicate the text embeddings for each generation, just to save time :)
+        bs_embed, seq_len, _ = text_embeddings.shape
+        text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
+        text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        return text_embeddings
+
+    def __call__(self, prompt, num_inference_steps=50, guidance_scale=7.5):
+        # encide the prompt into the text embeddings
+        text_embeddings = self.encode_prompt(prompt, self.device, 1, do_classifier_free_guidance=False)
+
+        # this is the initial noise sample
+        latents = torch.randn(
+            (text_embeddings.shape[0], self.unet.in_channels, self.unet.image_size, self.unet.image_size),
+            device=self.device
+        )
+
+        # setting up the scheduler
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        # iterate and generate
+        for t in self.scheduler.timesteps:
+            latents = self.scheduler.scale_model_input(latents, t)
+            latents = self.unet(latents, t, text_embeddings)["sample"]
+            latents = self.scheduler.step(latents, t, latents, guidance_scale=guidance_scale)["prev_sample"]
+
+        # decode latents into images
+        images = self.vae.decode(latents)
+        images = (images / 2 + 0.5).clamp(0, 1)
+
+        return images