Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 from huggingface_hub import InferenceClient
+import time
 
+# Initialize the HuggingFace Inference Client
 client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
 
 def format_prompt(message, history, system_prompt=""):
@@ -14,11 +16,7 @@ def format_prompt(message, history, system_prompt=""):
     return prompt
 
 def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
+    temperature = max(temperature, 1e-2)
     generate_kwargs = dict(
         temperature=temperature,
         max_new_tokens=max_new_tokens,
@@ -29,41 +27,48 @@ def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=
     )
 
     formatted_prompt = format_prompt(prompt, history, system_prompt)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-        st.session_state.generated_text = output
+    try:
+        stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream:
+            output += response.token.text
+        return output
+    except Exception as e:
+        st.error(f"Error generating text: {e}")
+        return ""
 
-    return output
+def main():
+    st.title("Waifu Character Generator")
 
-# Streamlit UI
-st.title("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", "")
 
-# 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", "")
+    # Advanced settings
+    with st.expander("Advanced Settings"):
+        temperature = st.slider("Temperature", 0.0, 1.0, 0.9, step=0.05)
+        max_new_tokens = st.slider("Max new tokens", 0, 8192, 512, step=64)
+        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)
 
-# Advanced settings
-temperature = st.slider("Temperature", 0.0, 1.0, 0.9, step=0.05)
-max_new_tokens = st.slider("Max new tokens", 0, 8192, 512, step=64)
-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 = ""
 
-# 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)
 
-# Generate button
-if st.button("Generate Waifu"):
-    history = []
-    prompt = f"Create a waifu character named {name} with {hair_color} hair, a {personality} personality, and wearing a {outfit_style}."
-    generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
-    
     # Display the generated character
-    st.subheader("Generated Waifu Character")
-    st.write(st.session_state.generated_text)
+    if st.session_state.generated_text:
+        st.subheader("Generated Waifu Character")
+        st.write(st.session_state.generated_text)
+
+if __name__ == "__main__":
+    main()