Update app.py

app.py

@@ -11,16 +11,17 @@ with open('style.css') as f:
 text_client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
 image_client = Client("Boboiazumi/animagine-xl-3.1")
 
-def format_prompt_for_description(name, hair_color, eye_color, height, age, personality, outfit_style, hobbies, favorite_food, favorite_quote, strengths, weaknesses, goals, background_story):
-    prompt = f"Create a waifu character named {name} with {hair_color} hair, {eye_color} eyes, a {personality} personality, and wearing a {outfit_style}. "
-    prompt += f"She is {height} tall and {age} years old. Her hobbies include {hobbies}. Her favorite food is {favorite_food}. "
-    prompt += f"Her favorite quote is: '{favorite_quote}'. She has the following strengths: {strengths}. Her weaknesses include {weaknesses}. "
-    prompt += f"Her goals are: {goals}. Here is her background story: {background_story}."
+def format_prompt_for_description(name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story, **kwargs):
+    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}."
+    for key, value in kwargs.items():
+        prompt += f" Her {key} is {value}."
     return prompt
 
-def format_prompt_for_image(name, hair_color, eye_color, personality, outfit_style, background_setting, pose):
-    prompt = f"Generate an image prompt for a waifu character named {name} with {hair_color} hair, {eye_color} eyes, a {personality} personality, and wearing a {outfit_style}. "
-    prompt += f"She is in a {background_setting} setting, posing in a {pose}."
+def format_prompt_for_image(name, hair_color, personality, outfit_style, **kwargs):
+    prompt = f"Generate an image prompt for a waifu character named {name} with {hair_color} hair, a {personality} personality, and wearing a {outfit_style}."
+    for key, value in kwargs.items():
+        prompt += f" She is in a {key} setting, posing in a {value}."
     return prompt
 
 def clean_generated_text(text):
@@ -93,74 +94,86 @@ def main():
     It's the best free waifu generator out there, providing everything you need to bring your waifu to life.
     """)
 
-    # User inputs
-    col1, col2 = st.columns(2)
-    with col1:
+    # Layout containers
+    left_col, right_col = st.columns(2)
+    
+    with left_col:
+        st.header("Waifu Settings")
+
+        # Basic settings
         name = st.text_input("Name of the Waifu")
         hair_color = st.selectbox("Hair Color", ["Blonde", "Brunette", "Red", "Black", "Blue", "Pink"])
-        eye_color = st.selectbox("Eye Color", ["Blue", "Green", "Brown", "Hazel", "Red", "Purple"])
-        height = st.text_input("Height (e.g., 5'6\")")
-        age = st.text_input("Age")
         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")
-        favorite_quote = st.text_input("Favorite Quote")
-        strengths = st.text_input("Strengths")
-        weaknesses = st.text_input("Weaknesses")
-        goals = st.text_input("Goals")
-        background_story = st.text_area("Background Story")
-        background_setting = st.selectbox("Background Setting", ["School", "City", "Fantasy World", "Beach", "Park"])
-        pose = st.selectbox("Pose", ["Standing", "Sitting", "Action Pose", "Lying Down"])
-        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)
-
-    # Initialize session state for generated text and image prompt
-    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 = ""
-    if "image_paths" not in st.session_state:
-        st.session_state.image_paths = []
-
-    # Generate button
-    if st.button("Generate Waifu"):
-        with st.spinner("Generating waifu character..."):
-            description_prompt = format_prompt_for_description(
-                name, hair_color, eye_color, height, age, personality, outfit_style, 
-                hobbies, favorite_food, favorite_quote, strengths, weaknesses, goals, background_story
-            )
-            image_prompt = format_prompt_for_image(name, hair_color, eye_color, personality, outfit_style, background_setting, pose)
-
-            # 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)
-            
-            # Generate image from image prompt
-            st.session_state.image_paths = generate_image(st.session_state.image_prompt)
-            
-            st.success("Waifu character generated!")
-
-    # Display the generated character and image prompt
-    if st.session_state.character_description:
-        st.subheader("Generated Waifu Character")
-        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 st.session_state.image_paths:
-        st.subheader("Generated Image")
-        for image_path in st.session_state.image_paths:
-            st.image(image_path, caption="Generated Waifu Image")
+
+        # Advanced settings
+        with st.expander("More Options"):
+            hobbies = st.text_input("Hobbies")
+            favorite_food = st.text_input("Favorite Food")
+            background_story = st.text_area("Background Story")
+            eye_color = st.selectbox("Eye Color", ["Blue", "Green", "Brown", "Hazel", "Red", "Purple"], index=0)
+            height = st.text_input("Height (e.g., 5'6\")", "")
+            age = st.text_input("Age", "")
+            favorite_quote = st.text_input("Favorite Quote", "")
+            strengths = st.text_input("Strengths", "")
+            weaknesses = st.text_input("Weaknesses", "")
+            goals = st.text_input("Goals", "")
+            background_setting = st.selectbox("Background Setting", ["School", "City", "Fantasy World", "Beach", "Park"], index=0)
+            pose = st.selectbox("Pose", ["Standing", "Sitting", "Action Pose", "Lying Down"], index=0)
+            system_prompt = st.text_input("Optional System Prompt", "")
+
+        # Advanced settings sliders
+        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)
+
+        # Initialize session state for generated text and image prompt
+        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 = ""
+        if "image_paths" not in st.session_state:
+            st.session_state.image_paths = []
+
+        # Generate button
+        if st.button("Generate Waifu"):
+            with st.spinner("Generating waifu character..."):
+                description_prompt = format_prompt_for_description(
+                    name, hair_color, personality, outfit_style, hobbies, favorite_food, background_story,
+                    eye_color=eye_color, height=height, age=age, favorite_quote=favorite_quote, strengths=strengths,
+                    weaknesses=weaknesses, goals=goals
+                )
+                image_prompt = format_prompt_for_image(
+                    name, hair_color, personality, outfit_style,
+                    background_setting=background_setting, pose=pose
+                )
+
+                # 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)
+                
+                # Generate image from image prompt
+                st.session_state.image_paths = generate_image(st.session_state.image_prompt)
+                
+                st.success("Waifu character generated!")
+
+    with right_col:
+        st.header("Generated Waifu")
+        # Display the generated character and image prompt
+        if st.session_state.character_description:
+            st.subheader("Generated Waifu Character")
+            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 st.session_state.image_paths:
+            st.subheader("Generated Image")
+            for image_path in st.session_state.image_paths:
+                st.image(image_path, caption="Generated Waifu Image")
 
 if __name__ == "__main__":
     main()