Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from huggingface_hub import InferenceClient
+import re
 
 # Load custom CSS
 with open('style.css') as f:
@@ -17,7 +18,12 @@ def format_prompt_for_image(name, hair_color, personality, outfit_style):
     prompt = f"Generate an image prompt for a waifu character named {name} with {hair_color} hair, a {personality} personality, and wearing a {outfit_style}."
     return prompt
 
-def generate_text(prompt, temperature=0.9, max_new_tokens=2512, top_p=0.95, repetition_penalty=1.0):
+def clean_generated_text(text):
+    # Remove any unwanted trailing tags or characters like </s>
+    clean_text = re.sub(r'</s>$', '', text).strip()
+    return clean_text
+
+def generate_text(prompt, temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
     temperature = max(temperature, 1e-2)
     generate_kwargs = dict(
         temperature=temperature,
@@ -32,7 +38,7 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=2512, top_p=0.95, repe
         output = ""
         for response in stream:
             output += response.token.text
-        return output
+        return clean_generated_text(output)
     except Exception as e:
         st.error(f"Error generating text: {e}")
         return ""