Update app.py

app.py

@@ -1,28 +1,44 @@
-import streamlit as st
-from gradio_client import Client
-
-def main():
-    st.title("Juggernaut-X-v10 Prediction")
-
-    # User input for the prompt
-    user_prompt = st.text_input("Enter your prompt:")
-
-    if st.button("Generate"):
-        with st.spinner('Generating...'):
-            # Call the Gradio API
-            client = Client("Walmart-the-bag/Juggernaut-X-v10")
-            result = client.predict(
-                prompt=user_prompt,
-                negative_prompt="easynegative, worst quality:1.4, bad composition:1.4, (bad anatomy:1.4, badly drawn hands:1.4, mutation, deformed, extra limb, missing limb, mutated limb, extra finger, missing finger), (deformed fingers:1.4, ugly fingers, ugly fingernails, deformed fingernails), watermarks, graphics, pixelation, blurry, artifacts, logo, (badly drawn face, badly drawn eyes, uneven eyes), imperfect proportions",
-                steps=50,
-                guidance_scale=7.5,
-                add_4k_masterpiece=False,
-                api_name="/predict"
-            )
-            # Display the result
-            st.success("Generated successfully!")
-            st.write(result)
-
-if __name__ == "__main__":
-    main()
-
+import streamlit as st
+from gradio_client import Client
+from PIL import Image
+import requests
+from io import BytesIO
+
+def main():
+    st.title("Juggernaut-X-v10 Prediction")
+
+    # User input for the prompt
+    user_prompt = st.text_input("Enter your prompt:")
+
+    if st.button("Generate"):
+        with st.spinner('Generating...'):
+            # Call the Gradio API
+            client = Client("Walmart-the-bag/Juggernaut-X-v10")
+            result = client.predict(
+                prompt=user_prompt,
+                negative_prompt="easynegative, worst quality:1.4, bad composition:1.4, (bad anatomy:1.4, badly drawn hands:1.4, mutation, deformed, extra limb, missing limb, mutated limb, extra finger, missing finger), (deformed fingers:1.4, ugly fingers, ugly fingernails, deformed fingernails), watermarks, graphics, pixelation, blurry, artifacts, logo, (badly drawn face, badly drawn eyes, uneven eyes), imperfect proportions",
+                steps=50,
+                guidance_scale=7.5,
+                add_4k_masterpiece=False,
+                api_name="/predict"
+            )
+            
+            # Display the result image
+            image_path = result['image']  # Adjust according to the actual key in the result dictionary
+            if image_path:
+                image_url = f"http://localhost:8501{image_path}"  # Adjust if needed
+                response = requests.get(image_url)
+                image = Image.open(BytesIO(response.content))
+                st.image(image, caption="Generated Image")
+
+                # Provide download button
+                image_bytes = BytesIO(response.content).getvalue()
+                st.download_button(
+                    label="Download Image",
+                    data=image_bytes,
+                    file_name="generated_image.webp",
+                    mime="image/webp"
+                )
+
+if __name__ == "__main__":
+    main()