From 857c04cfbafcc43cbf04e647d7cd1699390af32a Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Thu, 16 Feb 2023 23:59:53 +0200 Subject: [PATCH] [Pix2Pix0] Add utility function to get edit vector (#2383) uP --- .../pipeline_stable_diffusion_pix2pix_zero.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) 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 index af0a7e2a4b..42875d6690 100644 --- 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 @@ class StableDiffusionPix2PixZeroPipeline(DiffusionPipeline): """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__(