Update app.py

app.py

@@ -216,17 +216,19 @@ def generate_30(
     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
     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)
     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 along the sequence dimension (dim=1)
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b], dim=1)
+    print('catted shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
-
+    
     # 4. (Optional) Average the pooled embeddings
     prompt_embeds = prompt_embeds.mean(dim=1, keepdim=True)
+    print('averaged shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=1, keepdim=True)