Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 from huggingface_hub import InferenceClient
-from gradio_client import Client
 import re
 
 # Load custom CSS
@@ -8,8 +7,7 @@ with open('style.css') as f:
     st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
 
 # Initialize the HuggingFace Inference Client
-text_client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
-image_client = Client("cagliostrolab/animagine-xl-3.1")
+client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
 
 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}. "
@@ -25,7 +23,7 @@ def clean_generated_text(text):
     clean_text = re.sub(r'</s>$', '', text).strip()
     return clean_text
 
-def generate_text(client, 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,
@@ -45,32 +43,6 @@ def generate_text(client, prompt, temperature=0.9, max_new_tokens=512, top_p=0.9
         st.error(f"Error generating text: {e}")
         return ""
 
-def generate_image(prompt):
-    try:
-        result = image_client.predict(
-            prompt,  # Image prompt
-            "",  # Negative prompt
-            0,  # Seed
-            512,  # Width
-            512,  # Height
-            7.5,  # Guidance scale
-            25,  # Number of inference steps
-            'DPM++ 2M Karras',  # Sampler
-            '1024 x 1024',  # Aspect Ratio
-            'Anime',  # Style Preset
-            '(None)',  # Quality Tags Presets
-            True,  # Use Upscaler
-            0,  # Strength
-            1,  # Upscale by
-            True,  # Add Quality Tags
-            api_name="/run"
-        )
-        return result[0]['image']
-    except Exception as e:
-        st.error(f"Error generating image: {e}")
-        st.write("Full error details:", e)
-        return None
-
 def main():
     st.title("Enhanced Waifu Character Generator")
 
@@ -94,13 +66,11 @@ 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 and image prompt
+    # Initialize session state for 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 = ""
-    if "image_path" not in st.session_state:
-        st.session_state.image_path = ""
 
     # Generate button
     if st.button("Generate Waifu"):
@@ -109,13 +79,10 @@ def main():
             image_prompt = format_prompt_for_image(name, hair_color, personality, outfit_style)
 
             # Generate character description
-            st.session_state.character_description = generate_text(text_client, description_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            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(text_client, image_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
-            
-            # Generate image from image prompt
-            st.session_state.image_path = generate_image(st.session_state.image_prompt)
+            st.session_state.image_prompt = generate_text(image_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
             
             st.success("Waifu character generated!")
 
@@ -126,10 +93,6 @@ def main():
     if st.session_state.image_prompt:
         st.subheader("Image Prompt")
         st.write(st.session_state.image_prompt)
-    if st.session_state.image_path:
-        st.subheader("Generated Image")
-        st.image(st.session_state.image_path)
 
 if __name__ == "__main__":
     main()
-