use compel for prompt embeddings

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import spaces
 import torch
+from compel import Compel, ReturnedEmbeddingsType
 from diffusers import DiffusionPipeline
 
 
@@ -11,39 +12,23 @@ pipe = DiffusionPipeline.from_pretrained(
 )
 pipe.to('cuda')
 
-
-def build_embeddings(enhanced_prompt, negative_prompt=None):
-    max_length = pipe.tokenizer.model_max_length
-
-    input_ids = pipe.tokenizer(enhanced_prompt, return_tensors="pt").input_ids
-    input_ids = input_ids.to("cuda")
-
-    negative_ids = pipe.tokenizer(
-        negative_prompt or "",
-        truncation=False,
-        padding="max_length",
-        max_length=input_ids.shape[-1],
-        return_tensors="pt"
-    ).input_ids
-    negative_ids = negative_ids.to("cuda")
-
-    concat_embeds = []
-    neg_embeds = []
-    for i in range(0, input_ids.shape[-1], max_length):
-        concat_embeds.append(pipe.text_encoder(input_ids[:, i: i + max_length])[0])
-        neg_embeds.append(pipe.text_encoder(negative_ids[:, i: i + max_length])[0])
-
-    prompt_embeds = torch.cat(concat_embeds, dim=1)
-    negative_prompt_embeds = torch.cat(neg_embeds, dim=1)
-    return prompt_embeds, negative_prompt_embeds
+compel = Compel(
+  tokenizer=[pipe.tokenizer, pipe.tokenizer_2] ,
+  text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+  requires_pooled=[False, True]
+)
 
 
 @spaces.GPU
 def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples):
-    prompt_embeds, neg_prompt_embeds = build_embeddings(prompt, negative_prompt)
+    embeds, pooled = compel(prompt)
+    neg_embeds, neg_pooled = compel(negative_prompt)
     return pipe(
-        prompt_embeds=prompt_embeds,
-        negative_prompt_embeds=neg_prompt_embeds,
+        prompt_embeds=embeds,
+        pooled_prompt_embeds=pooled,
+        negative_prompt_embeds=neg_embeds,
+        negative_pooled_prompt_embeds=neg_pooled,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
         width=width,

requirements.txt

@@ -3,4 +3,5 @@ diffusers
 invisible_watermark
 torch
 transformers
-xformers
+xformers
+compel