Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from huggingface_hub import InferenceClient
+from gradio_client import Client
 import re
 
 # Load custom CSS
@@ -7,7 +8,8 @@ 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")
+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, 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}. "
@@ -34,7 +36,7 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=512, top_p=0.95, repet
         seed=42,
     )
     try:
-        stream = client.text_generation(prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        stream = text_client.text_generation(prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
         output = ""
         for response in stream:
             output += response.token.text
@@ -43,6 +45,35 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=512, top_p=0.95, repet
         st.error(f"Error generating text: {e}")
         return ""
 
+def generate_image(prompt):
+    try:
+        result = image_client.predict(
+            prompt=prompt,
+            negative_prompt="",
+            seed=0,
+            custom_width=1024,
+            custom_height=1024,
+            guidance_scale=7.0,
+            num_inference_steps=28,
+            sampler="Euler a",
+            aspect_ratio_selector="896 x 1152",
+            style_selector="(None)",
+            quality_selector="Standard v3.1",
+            use_upscaler=False,
+            upscaler_strength=0.55,
+            upscale_by=1.5,
+            add_quality_tags=True,
+            isImg2Img=False,
+            img_path=None,
+            img2img_strength=0.65,
+            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")
 
@@ -66,11 +97,13 @@ 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
+    # 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_path" not in st.session_state:
+        st.session_state.image_path = ""
 
     # Generate button
     if st.button("Generate Waifu"):
@@ -84,6 +117,9 @@ def main():
             # 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_path = generate_image(st.session_state.image_prompt)
+            
             st.success("Waifu character generated!")
 
     # Display the generated character and image prompt
@@ -93,6 +129,9 @@ 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, caption="Generated Waifu Image")
 
 if __name__ == "__main__":
     main()