Update app.py

app.py

@@ -193,6 +193,8 @@ def generate_30(
     pipe.text_encoder=text_encoder.to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_2=text_encoder_2.to(device=device, dtype=torch.bfloat16)
 
+    pooled_prompt_embeds_list=[]
+    prompt_embeds_list=[]
     text_inputs1 = pipe.tokenizer(
                     prompt,
                     padding="max_length",
@@ -210,11 +212,15 @@ def generate_30(
     )
     text_input_ids2 = text_inputs2.input_ids
     prompt_embedsa = pipe.text_encoder(text_input_ids1.to(device), output_hidden_states=True)
+    pooled_prompt_embeds_list.append(prompt_embedsa)
     prompt_embedsa = prompt_embedsa.hidden_states[-2]
+    prompt_embeds_list.append(prompt_embedsa[0])
     prompt_embedsb = pipe.text_encoder(text_input_ids2.to(device), output_hidden_states=True)
+    pooled_prompt_embeds_list.append(prompt_embedsb[0])
     prompt_embedsb = prompt_embedsb.hidden_states[-2]
-    prompt_embeds = torch.cat([prompt_embedsa,prompt_embedsb]).mean(dim=-1)
-    pooled_prompt_embeds = prompt_embeds[0]
+    prompt_embeds_list.append(prompt_embedsb)
+    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=-1)
+    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=-1)
     
     options = {
         #"prompt": prompt,
@@ -262,6 +268,8 @@ def generate_60(
     pipe.text_encoder=text_encoder.to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_2=text_encoder_2.to(device=device, dtype=torch.bfloat16)
 
+    pooled_prompt_embeds_list=[]
+    prompt_embeds_list=[]
     text_inputs1 = pipe.tokenizer(
                     prompt,
                     padding="max_length",
@@ -279,11 +287,15 @@ def generate_60(
     )
     text_input_ids2 = text_inputs2.input_ids
     prompt_embedsa = pipe.text_encoder(text_input_ids1.to(device), output_hidden_states=True)
+    pooled_prompt_embeds_list.append(prompt_embedsa)
     prompt_embedsa = prompt_embedsa.hidden_states[-2]
+    prompt_embeds_list.append(prompt_embedsa[0])
     prompt_embedsb = pipe.text_encoder(text_input_ids2.to(device), output_hidden_states=True)
+    pooled_prompt_embeds_list.append(prompt_embedsb[0])
     prompt_embedsb = prompt_embedsb.hidden_states[-2]
-    prompt_embeds = torch.cat([prompt_embedsa,prompt_embedsb]).mean(dim=-1)
-    pooled_prompt_embeds = prompt_embeds[0]
+    prompt_embeds_list.append(prompt_embedsb)
+    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=-1)
+    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=-1)
     
     options = {
         #"prompt": prompt,
@@ -331,6 +343,8 @@ def generate_90(
     pipe.text_encoder=text_encoder.to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_2=text_encoder_2.to(device=device, dtype=torch.bfloat16)
 
+    pooled_prompt_embeds_list=[]
+    prompt_embeds_list=[]
     text_inputs1 = pipe.tokenizer(
                     prompt,
                     padding="max_length",
@@ -348,11 +362,15 @@ def generate_90(
     )
     text_input_ids2 = text_inputs2.input_ids
     prompt_embedsa = pipe.text_encoder(text_input_ids1.to(device), output_hidden_states=True)
+    pooled_prompt_embeds_list.append(prompt_embedsa)
     prompt_embedsa = prompt_embedsa.hidden_states[-2]
+    prompt_embeds_list.append(prompt_embedsa[0])
     prompt_embedsb = pipe.text_encoder(text_input_ids2.to(device), output_hidden_states=True)
+    pooled_prompt_embeds_list.append(prompt_embedsb[0])
     prompt_embedsb = prompt_embedsb.hidden_states[-2]
-    prompt_embeds = torch.cat([prompt_embedsa,prompt_embedsb]).mean(dim=-1)
-    pooled_prompt_embeds = prompt_embeds[0]
+    prompt_embeds_list.append(prompt_embedsb)
+    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=-1)
+    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=-1)
     
     options = {
         #"prompt": prompt,