Update app.py

app.py

@@ -211,19 +211,24 @@ def generate_30(
                     return_tensors="pt",
     )
     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[0])
-    prompt_embedsa = prompt_embedsa.hidden_states[-2]
-    print('text_encoder shape: ',prompt_embedsa.shape)
-    prompt_embeds_list.append(prompt_embedsa)
-    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_list.append(prompt_embedsb)
-    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=1, keepdim=True)
-    print('catted shape: ',prompt_embeds.shape)
-    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=1, keepdim=True)
-    
+
+    # 2. Encode with the two text encoders
+    prompt_embeds_a = pipe.text_encoder(text_input_ids1, 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
+
+    prompt_embeds_b = pipe.text_encoder_2(text_input_ids2, 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)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
+
+    # 4. (Optional) Average the pooled embeddings
+    pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=1, keepdim=True)
+
+
     options = {
         #"prompt": prompt,
         "prompt_embeds": prompt_embeds,
@@ -288,19 +293,24 @@ def generate_60(
                     return_tensors="pt",
     )
     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[0])
-    prompt_embedsa = prompt_embedsa.hidden_states[-2]
-    print('text_encoder shape: ',prompt_embedsa.shape)
-    prompt_embeds_list.append(prompt_embedsa)
-    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_list.append(prompt_embedsb)
-    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=1, keepdim=True)
-    print('catted shape: ',prompt_embeds.shape)
-    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=1, keepdim=True)
-    
+
+    # 2. Encode with the two text encoders
+    prompt_embeds_a = pipe.text_encoder(text_input_ids1, 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
+
+    prompt_embeds_b = pipe.text_encoder_2(text_input_ids2, 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)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
+
+    # 4. (Optional) Average the pooled embeddings
+    pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=1, keepdim=True)
+
+
     options = {
         #"prompt": prompt,
         "prompt_embeds": prompt_embeds,
@@ -365,19 +375,24 @@ def generate_90(
                     return_tensors="pt",
     )
     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[0])
-    prompt_embedsa = prompt_embedsa.hidden_states[-2]
-    print('text_encoder shape: ',prompt_embedsa.shape)
-    prompt_embeds_list.append(prompt_embedsa)
-    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_list.append(prompt_embedsb)
-    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=1, keepdim=True)
-    print('catted shape: ',prompt_embeds.shape)
-    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=1, keepdim=True)
-    
+
+    # 2. Encode with the two text encoders
+    prompt_embeds_a = pipe.text_encoder(text_input_ids1, 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
+
+    prompt_embeds_b = pipe.text_encoder_2(text_input_ids2, 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)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
+
+    # 4. (Optional) Average the pooled embeddings
+    pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=1, keepdim=True)
+
+
     options = {
         #"prompt": prompt,
         "prompt_embeds": prompt_embeds,