Update app.py

app.py

@@ -8,14 +8,16 @@ with open('style.css') as f:
 # Initialize the HuggingFace Inference Client
 client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
 
-def format_prompt(name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story, system_prompt=""):
+def format_prompt_for_description(name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story):
     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] " + prompt
     return prompt
 
-def generate_text(prompt, temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
+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):
     temperature = max(temperature, 1e-2)
     generate_kwargs = dict(
         temperature=temperature,
@@ -59,20 +61,32 @@ def main():
         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 = ""
+    if "character_description" not in st.session_state:
+        st.session_state.character_description = ""
+    if "image_prompt" not in st.session_state:
+        st.session_state.image_prompt = ""
 
     # Generate button
     if st.button("Generate Waifu"):
         with st.spinner("Generating waifu character..."):
-            prompt = format_prompt(name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story, system_prompt)
-            st.session_state.generated_text = generate_text(prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            description_prompt = format_prompt_for_description(name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story)
+            image_prompt = format_prompt_for_image(name, hair_color, personality, outfit_style)
+
+            # Generate character description
+            st.session_state.character_description = generate_text(description_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            
+            # Generate image prompt
+            st.session_state.image_prompt = generate_text(image_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            
             st.success("Waifu character generated!")
 
-    # Display the generated character
-    if st.session_state.generated_text:
+    # Display the generated character and image prompt
+    if st.session_state.character_description:
         st.subheader("Generated Waifu Character")
-        st.write(st.session_state.generated_text)
+        st.write(st.session_state.character_description)
+    if st.session_state.image_prompt:
+        st.subheader("Image Prompt")
+        st.write(st.session_state.image_prompt)
 
 if __name__ == "__main__":
     main()