Update app.py

app.py

@@ -195,6 +195,7 @@ def generate_30(
 
     pooled_prompt_embeds_list=[]
     prompt_embeds_list=[]
+    
     text_inputs1 = pipe.tokenizer(
                     prompt,
                     padding="max_length",
@@ -202,7 +203,9 @@ def generate_30(
                     truncation=True,
                     return_tensors="pt",
     )
+    
     text_input_ids1 = text_inputs1.input_ids
+    
     text_inputs2 = pipe.tokenizer(
                     prompt2,
                     padding="max_length",
@@ -210,8 +213,28 @@ def generate_30(
                     truncation=True,
                     return_tensors="pt",
     )
+    
     text_input_ids2 = text_inputs2.input_ids
 
+    text_inputs1b = pipe.tokenizer_2(
+                    prompt,
+                    padding="max_length",
+                    max_length=77,
+                    truncation=True,
+                    return_tensors="pt",
+    )
+
+    text_input_ids1b = text_inputs1b.input_ids
+
+    text_inputs2b = pipe.tokenizer_2(
+                    prompt2,
+                    padding="max_length",
+                    max_length=77,
+                    truncation=True,
+                    return_tensors="pt",
+    )
+    text_input_ids2b = text_inputs2b.input_ids
+
     # 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
@@ -221,6 +244,14 @@ def generate_30(
     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
 
+    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)
+    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
+
     # 3. Concatenate the embeddings 
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
@@ -231,6 +262,15 @@ def generate_30(
     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)
+    # 4. (Optional) Average the pooled embeddings
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds2,dim=0)
+    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
+    
     options = {
         #"prompt": prompt,
         "prompt_embeds": prompt_embeds,
@@ -279,6 +319,7 @@ def generate_60(
 
     pooled_prompt_embeds_list=[]
     prompt_embeds_list=[]
+    
     text_inputs1 = pipe.tokenizer(
                     prompt,
                     padding="max_length",
@@ -286,7 +327,9 @@ def generate_60(
                     truncation=True,
                     return_tensors="pt",
     )
+    
     text_input_ids1 = text_inputs1.input_ids
+    
     text_inputs2 = pipe.tokenizer(
                     prompt2,
                     padding="max_length",
@@ -294,8 +337,28 @@ def generate_60(
                     truncation=True,
                     return_tensors="pt",
     )
+    
     text_input_ids2 = text_inputs2.input_ids
 
+    text_inputs1b = pipe.tokenizer_2(
+                    prompt,
+                    padding="max_length",
+                    max_length=77,
+                    truncation=True,
+                    return_tensors="pt",
+    )
+
+    text_input_ids1b = text_inputs1b.input_ids
+
+    text_inputs2b = pipe.tokenizer_2(
+                    prompt2,
+                    padding="max_length",
+                    max_length=77,
+                    truncation=True,
+                    return_tensors="pt",
+    )
+    text_input_ids2b = text_inputs2b.input_ids
+
     # 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
@@ -305,6 +368,14 @@ def generate_60(
     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
 
+    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)
+    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
+
     # 3. Concatenate the embeddings 
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
@@ -315,7 +386,15 @@ def generate_60(
     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)
+    # 4. (Optional) Average the pooled embeddings
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds2,dim=0)
+    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
+    
     options = {
         #"prompt": prompt,
         "prompt_embeds": prompt_embeds,
@@ -364,6 +443,7 @@ def generate_90(
 
     pooled_prompt_embeds_list=[]
     prompt_embeds_list=[]
+    
     text_inputs1 = pipe.tokenizer(
                     prompt,
                     padding="max_length",
@@ -371,7 +451,9 @@ def generate_90(
                     truncation=True,
                     return_tensors="pt",
     )
+    
     text_input_ids1 = text_inputs1.input_ids
+    
     text_inputs2 = pipe.tokenizer(
                     prompt2,
                     padding="max_length",
@@ -379,8 +461,28 @@ def generate_90(
                     truncation=True,
                     return_tensors="pt",
     )
+    
     text_input_ids2 = text_inputs2.input_ids
 
+    text_inputs1b = pipe.tokenizer_2(
+                    prompt,
+                    padding="max_length",
+                    max_length=77,
+                    truncation=True,
+                    return_tensors="pt",
+    )
+
+    text_input_ids1b = text_inputs1b.input_ids
+
+    text_inputs2b = pipe.tokenizer_2(
+                    prompt2,
+                    padding="max_length",
+                    max_length=77,
+                    truncation=True,
+                    return_tensors="pt",
+    )
+    text_input_ids2b = text_inputs2b.input_ids
+
     # 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
@@ -389,7 +491,15 @@ 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
-    
+
+    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)
+    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
+
     # 3. Concatenate the embeddings 
     prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
@@ -400,6 +510,15 @@ def generate_90(
     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)
+    # 4. (Optional) Average the pooled embeddings
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds2,dim=0)
+    print('pooled averaged shape: ', pooled_prompt_embeds.shape)
+    
     options = {
         #"prompt": prompt,
         "prompt_embeds": prompt_embeds,