Update app.py

app.py

@@ -228,7 +228,8 @@ def generate_30(
     # 4. (Optional) Average the pooled embeddings
     prompt_embeds = torch.mean(prompt_embeds,dim=0,keepdim=True)
     print('averaged shape: ', prompt_embeds.shape)
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
+    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
 
     options = {
         #"prompt": prompt,
@@ -311,7 +312,8 @@ def generate_60(
     # 4. (Optional) Average the pooled embeddings
     prompt_embeds = torch.mean(prompt_embeds,dim=0,keepdim=True)
     print('averaged shape: ', prompt_embeds.shape)
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
+    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
 
 
     options = {
@@ -387,7 +389,7 @@ def generate_90(
     prompt_embeds_b = pipe.text_encoder(text_input_ids2.to(torch.device('cuda')), output_hidden_states=True)
     pooled_prompt_embeds_b = prompt_embeds_b[0]  # Pooled output from encoder 2
     prompt_embeds_b = prompt_embeds_b.hidden_states[-2]  # Penultimate hidden state from encoder 2
-
+    
     # 3. Concatenate the embeddings 
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
@@ -395,7 +397,8 @@ def generate_90(
     # 4. (Optional) Average the pooled embeddings
     prompt_embeds = torch.mean(prompt_embeds,dim=0,keepdim=True)
     print('averaged shape: ', prompt_embeds.shape)
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
+    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
 
     options = {
         #"prompt": prompt,