Update app.py

app.py

@@ -305,6 +305,10 @@ def expand_prompt(prompt):
         )
         enhanced_prompt = txt_tokenizer.decode(outputs[0], skip_special_tokens=True)
         enhanced_prompt = filter_text(enhanced_prompt,prompt)
+        enhanced_prompt = filter_text(enhanced_prompt,user_prompt_rewrite)
+        enhanced_prompt = filter_text(enhanced_prompt,system_prompt_rewrite)
+        print('-- filtered prompt --')
+        print(enhanced_prompt)
         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")
@@ -319,10 +323,9 @@ def expand_prompt(prompt):
         )
         # Use the encoded tensor 'text_inputs' here
         enhanced_prompt_2 = txt_tokenizer.decode(outputs_2[0], skip_special_tokens=True)
-        print('-- generated prompt --')
         enhanced_prompt_2 = filter_text(enhanced_prompt_2,prompt)
-        print('-- filtered prompt --')
-        print(enhanced_prompt)
+        enhanced_prompt_2 = filter_text(enhanced_prompt_2,user_prompt_rewrite_2)
+        enhanced_prompt_2 = filter_text(enhanced_prompt_2,system_prompt_rewrite)
         print('-- filtered prompt 2 --')
         print(enhanced_prompt_2)
         enh_prompt=[enhanced_prompt,enhanced_prompt_2]
@@ -426,7 +429,7 @@ def generate_30(
         expanded = expand_prompt(caption)
         expanded_1 = expanded[0]
         expanded_2 = expanded[1]
-        new_prompt = prompt+' '+expanded_1
+        new_prompt = prompt+' '+expanded_1+' '+expanded_2
         print("-- ------------ --")
         print("-- FINAL PROMPT --")
         print(new_prompt)
@@ -577,7 +580,7 @@ def generate_60(
         expanded = expand_prompt(caption)
         expanded_1 = expanded[0]
         expanded_2 = expanded[1]
-        new_prompt = prompt+' '+expanded_1
+        new_prompt = prompt+' '+expanded_1+' '+expanded_2
         print("-- ------------ --")
         print("-- FINAL PROMPT --")
         print(new_prompt)
@@ -728,7 +731,7 @@ def generate_90(
         expanded = expand_prompt(caption)
         expanded_1 = expanded[0]
         expanded_2 = expanded[1]
-        new_prompt = prompt+' '+expanded_1
+        new_prompt = prompt+' '+expanded_1+' '+expanded_2
         print("-- ------------ --")
         print("-- FINAL PROMPT --")
         print(new_prompt)