Update app.py

app.py

@@ -258,7 +258,7 @@ def generate_30(
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
-    print('catted pooled shape: ', prompt_embeds.shape)
+    print('catted pooled shape: ', pooled_prompt_embeds.shape)
     pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
     print('meaned pooled shape: ', pooled_prompt_embeds.shape)
 
@@ -271,7 +271,7 @@ def generate_30(
     print('catted shape2: ', prompt_embeds2.shape)
     pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds_a2, pooled_prompt_embeds_b2])
     print('catted pooled shape 2: ', pooled_prompt_embeds2.shape)
-    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=0,keepdim=True)
+    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=0)
     print('pooled meaned shape 2: ', pooled_prompt_embeds2.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2])
     print('catted combined meaned pooled shape: ', pooled_prompt_embeds.shape)
@@ -367,6 +367,7 @@ def generate_60(
     # 2. Encode with the two text encoders
     prompt_embeds_a = pipe.text_encoder(text_input_ids1.to(torch.device('cuda')), output_hidden_states=True)
     pooled_prompt_embeds_a = prompt_embeds_a[0]  # Pooled output from encoder 1
+    print('pooled shape 1: ', pooled_prompt_embeds_a.shape)
     prompt_embeds_a = prompt_embeds_a.hidden_states[-2]  # Penultimate hidden state from encoder 1
     print('encoder shape: ', prompt_embeds_a.shape)
     prompt_embeds_b = pipe.text_encoder(text_input_ids2.to(torch.device('cuda')), output_hidden_states=True)
@@ -375,7 +376,7 @@ def generate_60(
 
     prompt_embeds_a2 = pipe.text_encoder_2(text_input_ids1b.to(torch.device('cuda')), output_hidden_states=True)
     pooled_prompt_embeds_a2 = prompt_embeds_a2[0]  # Pooled output from encoder 1
-    print('pooled shape: ', pooled_prompt_embeds_a2.shape)
+    print('pooled shape 2: ', pooled_prompt_embeds_a2.shape)
     prompt_embeds_a2 = prompt_embeds_a2.hidden_states[-2]  # Penultimate hidden state from encoder 1
     print('encoder shape2: ', prompt_embeds_a2.shape)
     prompt_embeds_b2 = pipe.text_encoder_2(text_input_ids2b.to(torch.device('cuda')), output_hidden_states=True)
@@ -386,7 +387,7 @@ def generate_60(
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
-    print('catted pooled shape: ', prompt_embeds.shape)
+    print('catted pooled shape: ', pooled_prompt_embeds.shape)
     pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
     print('meaned pooled shape: ', pooled_prompt_embeds.shape)
 
@@ -495,6 +496,7 @@ def generate_90(
     # 2. Encode with the two text encoders
     prompt_embeds_a = pipe.text_encoder(text_input_ids1.to(torch.device('cuda')), output_hidden_states=True)
     pooled_prompt_embeds_a = prompt_embeds_a[0]  # Pooled output from encoder 1
+    print('pooled shape 1: ', pooled_prompt_embeds_a.shape)
     prompt_embeds_a = prompt_embeds_a.hidden_states[-2]  # Penultimate hidden state from encoder 1
     print('encoder shape: ', prompt_embeds_a.shape)
     prompt_embeds_b = pipe.text_encoder(text_input_ids2.to(torch.device('cuda')), output_hidden_states=True)
@@ -503,7 +505,7 @@ def generate_90(
 
     prompt_embeds_a2 = pipe.text_encoder_2(text_input_ids1b.to(torch.device('cuda')), output_hidden_states=True)
     pooled_prompt_embeds_a2 = prompt_embeds_a2[0]  # Pooled output from encoder 1
-    print('pooled shape: ', pooled_prompt_embeds_a2.shape)
+    print('pooled shape 2: ', pooled_prompt_embeds_a2.shape)
     prompt_embeds_a2 = prompt_embeds_a2.hidden_states[-2]  # Penultimate hidden state from encoder 1
     print('encoder shape2: ', prompt_embeds_a2.shape)
     prompt_embeds_b2 = pipe.text_encoder_2(text_input_ids2b.to(torch.device('cuda')), output_hidden_states=True)
@@ -514,7 +516,7 @@ def generate_90(
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
-    print('catted pooled shape: ', prompt_embeds.shape)
+    print('catted pooled shape: ', pooled_prompt_embeds.shape)
     pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
     print('meaned pooled shape: ', pooled_prompt_embeds.shape)