Update app.py

app.py

@@ -247,7 +247,7 @@ def generate_30(
     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
     prompt_embeds_a2 = prompt_embeds_a2.hidden_states[-2]  # Penultimate hidden state from encoder 1
-    print('encoder shape: ', prompt_embeds_a2.shape)
+    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)
     pooled_prompt_embeds_b2 = prompt_embeds_b2[0]  # Pooled output from encoder 2
     prompt_embeds_b2 = prompt_embeds_b2.hidden_states[-2]  # Penultimate hidden state from encoder 2
@@ -259,16 +259,14 @@ 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)
-    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
 
     # 3. Concatenate the text_encoder_2 embeddings 
     prompt_embeds2 = torch.cat([prompt_embeds_a2, prompt_embeds_b2])
     print('catted shape2: ', prompt_embeds.shape)
     pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds_a2, pooled_prompt_embeds_b2])
-    pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2], dim=2)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2])
     # 4. (Optional) Average the pooled embeddings
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds2,dim=0)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
     print('pooled averaged shape: ', pooled_prompt_embeds.shape)
     
     options = {
@@ -376,6 +374,7 @@ def generate_60(
     pooled_prompt_embeds_b2 = prompt_embeds_b2[0]  # Pooled output from encoder 2
     prompt_embeds_b2 = prompt_embeds_b2.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)
@@ -383,16 +382,14 @@ 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)
-    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
 
     # 3. Concatenate the text_encoder_2 embeddings 
     prompt_embeds2 = torch.cat([prompt_embeds_a2, prompt_embeds_b2])
     print('catted shape2: ', prompt_embeds.shape)
     pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds_a2, pooled_prompt_embeds_b2])
-    pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2], dim=2)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2])
     # 4. (Optional) Average the pooled embeddings
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds2,dim=0)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
     print('pooled averaged shape: ', pooled_prompt_embeds.shape)
     
     options = {
@@ -500,6 +497,7 @@ def generate_90(
     pooled_prompt_embeds_b2 = prompt_embeds_b2[0]  # Pooled output from encoder 2
     prompt_embeds_b2 = prompt_embeds_b2.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)
@@ -507,16 +505,14 @@ 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)
-    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
 
     # 3. Concatenate the text_encoder_2 embeddings 
     prompt_embeds2 = torch.cat([prompt_embeds_a2, prompt_embeds_b2])
     print('catted shape2: ', prompt_embeds.shape)
     pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds_a2, pooled_prompt_embeds_b2])
-    pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2], dim=2)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2])
     # 4. (Optional) Average the pooled embeddings
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds2,dim=0)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
     print('pooled averaged shape: ', pooled_prompt_embeds.shape)
     
     options = {