Update app.py

app.py

@@ -279,37 +279,37 @@ def expand_prompt(prompt):
             "Rephrase this scene to have more elaborate details: "
         )
         input_text = f"{system_prompt_rewrite} {user_prompt_rewrite} {prompt}"
-        input_text_2 = f"{system_prompt_rewrite} {user_prompt_rewrite_2} {prompt}"
         print("-- got prompt --")
         # Encode the input text and include the attention mask
         encoded_inputs = txt_tokenizer(input_text, return_tensors="pt", return_attention_mask=True).to("cuda:0")
-        encoded_inputs_2 = txt_tokenizer(input_text_2, return_tensors="pt", return_attention_mask=True).to("cuda:0")
         # Ensure all values are on the correct device
         input_ids = encoded_inputs["input_ids"].to("cuda:0")
-        input_ids_2 = encoded_inputs_2["input_ids"].to("cuda:0")
         attention_mask = encoded_inputs["attention_mask"].to("cuda:0")
-        attention_mask_2 = encoded_inputs_2["attention_mask"].to("cuda:0")
         print("-- tokenize prompt --")
           # Google T5
         #input_ids = txt_tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
         outputs = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            max_new_tokens=512,
+            max_new_tokens=1024,
             temperature=0.2,
             top_p=0.9,
             do_sample=True,
         )
+        enhanced_prompt = txt_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        input_text_2 = f"{system_prompt_rewrite} {user_prompt_rewrite_2} {enhanced_prompt}"
+        encoded_inputs_2 = txt_tokenizer(input_text_2, return_tensors="pt", return_attention_mask=True).to("cuda:0")
+        input_ids_2 = encoded_inputs_2["input_ids"].to("cuda:0")
+        attention_mask_2 = encoded_inputs_2["attention_mask"].to("cuda:0")
         outputs_2 = model.generate(
             input_ids=input_ids_2,
             attention_mask=attention_mask_2,
-            max_new_tokens=65,
+            max_new_tokens=1024,
             temperature=0.2,
             top_p=0.9,
             do_sample=True,
         )
         # Use the encoded tensor 'text_inputs' here
-        enhanced_prompt = txt_tokenizer.decode(outputs[0], skip_special_tokens=True)
         enhanced_prompt_2 = txt_tokenizer.decode(outputs_2[0], skip_special_tokens=True)
         print('-- generated prompt --')
         enhanced_prompt = filter_text(enhanced_prompt,prompt)
@@ -404,7 +404,7 @@ def generate_30(
 
         expand_prompt(prompt)
         expand_prompt(caption)
-        expand_prompt(caption_2)
+        expanded = expand_prompt(caption_2)
       
         print('-- generating image --')
         sd_image = ip_model.generate(
@@ -414,6 +414,7 @@ def generate_30(
                 pil_image_4=sd_image_d,
                 pil_image_5=sd_image_e,
                 prompt=prompt,
+                prompt_2=expanded,
                 negative_prompt=negative_prompt,
                 text_scale=text_scale,
                 ip_scale=ip_scale,