[Pix2Pix0] Add utility function to get edit vector (open-mmlab#2383)

patrickvonplaten · web-flow · commit 857c04cfbafc · 2023-02-16T22:59:53.000+01:00
uP
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py
@@ -542,6 +542,20 @@ def construct_direction(self, embs_source: torch.Tensor, embs_target: torch.Tens
         """Constructs the edit direction to steer the image generation process semantically."""
         return (embs_target.mean(0) - embs_source.mean(0)).unsqueeze(0)
 
+    @torch.no_grad()
+    def get_embeds(self, prompt: List[str]) -> torch.FloatTensor:
+        input_ids = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        ).input_ids
+
+        embeds = self.text_encoder(input_ids)[0]
+
+        return embeds.mean(0)[None]
+
     @torch.no_grad()
     @replace_example_docstring(EXAMPLE_DOC_STRING)
     def __call__(
