Update app.py

app.py

@@ -2,20 +2,21 @@ import streamlit as st
 from huggingface_hub import InferenceClient
 import time
 
+# Load custom CSS
+with open('style.css') as f:
+    st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
+
 # Initialize the HuggingFace Inference Client
 client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
 
-def format_prompt(message, history, system_prompt=""):
-    prompt = "<s>"
+def format_prompt(name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story, system_prompt=""):
+    prompt = f"Create a waifu character named {name} with {hair_color} hair, a {personality} personality, and wearing a {outfit_style}. "
+    prompt += f"Her hobbies include {hobbies}. Her favorite food is {favorite_food}. Here is her background story: {background_story}."
     if system_prompt:
-        prompt += f"[SYS] {system_prompt} [/SYS] "
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
+        prompt = f"[SYS] {system_prompt} [/SYS] " + prompt
     return prompt
 
-def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
+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,
@@ -25,10 +26,8 @@ def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=
         do_sample=True,
         seed=42,
     )
-
-    formatted_prompt = format_prompt(prompt, history, system_prompt)
     try:
-        stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        stream = client.text_generation(prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
         output = ""
         for response in stream:
             output += response.token.text
@@ -38,14 +37,20 @@ def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=
         return ""
 
 def main():
-    st.title("Waifu Character Generator")
+    st.title("Enhanced Waifu Character Generator")
 
     # User inputs
-    name = st.text_input("Name of the Waifu")
-    hair_color = st.selectbox("Hair Color", ["Blonde", "Brunette", "Red", "Black", "Blue", "Pink"])
-    personality = st.selectbox("Personality", ["Tsundere", "Yandere", "Kuudere", "Dandere", "Genki", "Normal"])
-    outfit_style = st.selectbox("Outfit Style", ["School Uniform", "Maid Outfit", "Casual", "Kimono", "Gothic Lolita"])
-    system_prompt = st.text_input("Optional System Prompt", "")
+    col1, col2 = st.columns(2)
+    with col1:
+        name = st.text_input("Name of the Waifu")
+        hair_color = st.selectbox("Hair Color", ["Blonde", "Brunette", "Red", "Black", "Blue", "Pink"])
+        personality = st.selectbox("Personality", ["Tsundere", "Yandere", "Kuudere", "Dandere", "Genki", "Normal"])
+        outfit_style = st.selectbox("Outfit Style", ["School Uniform", "Maid Outfit", "Casual", "Kimono", "Gothic Lolita"])
+    with col2:
+        hobbies = st.text_input("Hobbies")
+        favorite_food = st.text_input("Favorite Food")
+        background_story = st.text_area("Background Story")
+        system_prompt = st.text_input("Optional System Prompt", "")
 
     # Advanced settings
     with st.expander("Advanced Settings"):
@@ -54,21 +59,17 @@ def main():
         top_p = st.slider("Top-p (nucleus sampling)", 0.0, 1.0, 0.95, step=0.05)
         repetition_penalty = st.slider("Repetition penalty", 1.0, 2.0, 1.0, step=0.05)
 
-    # Initialize session state for generated text
-    if "generated_text" not in st.session_state:
-        st.session_state.generated_text = ""
-
     # Generate button
     if st.button("Generate Waifu"):
         with st.spinner("Generating waifu character..."):
-            history = []
-            prompt = f"Create a waifu character named {name} with {hair_color} hair, a {personality} personality, and wearing a {outfit_style}."
-            st.session_state.generated_text = generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            prompt = format_prompt(name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story, system_prompt)
+            generated_text = generate_text(prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            st.success("Waifu character generated!")
 
     # Display the generated character
-    if st.session_state.generated_text:
+    if generated_text:
         st.subheader("Generated Waifu Character")
-        st.write(st.session_state.generated_text)
+        st.write(generated_text)
 
 if __name__ == "__main__":
     main()