Update app.py

app.py

@@ -81,7 +81,7 @@ tokenizer.tokenizer_legacy=False
 model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map='balanced')
 #model = torch.compile(model)
 
-def filter_text(text):
+def filter_text(text,phraseC):
   """Filters out the text up to and including 'Rewritten Prompt:'."""
   phrase = "Rewritten Prompt:"
   phraseB = "rewritten text:"
@@ -94,6 +94,7 @@ def filter_text(text):
         match = re.search(pattern, filtered_text, flags=re.DOTALL)
         if match:
           filtered_text = match.group(2)
+          filtered_text = re.sub(phraseC, "", filtered_text, flags=re.DOTALL)  # Replaces the matched pattern with an empty string
           return filtered_text
         else:
           return filtered_text
@@ -148,7 +149,7 @@ def infer(
     enhanced_prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
     print('-- generated prompt --')
     print(enhanced_prompt)
-    enhanced_prompt = filter_text(enhanced_prompt)
+    enhanced_prompt = filter_text(enhanced_prompt,prompt)
     print('-- filtered prompt --')
     print(enhanced_prompt)
     print('-- generating image --')