Update src/streamlit_app.py

src/streamlit_app.py

@@ -211,12 +211,20 @@ st.markdown("---")
 
 # ---------------------- GENERATE IMAGE ----------------------
 st.subheader("🖼️ Step 6: Generate Image")
+
+if "generated_image" not in st.session_state:
+    st.session_state.generated_image = None
+
 if st.button("Generate Image from Story"):
     resp = requests.post(f"{API_URL}/generate-image")
     if resp.ok:
-        img_url = resp.json()['response']
-        # st.image("image.png", caption="Generated Image from Story")
-        st.session_state.generated_image = "image.png"
+        # Use static image URL (if backend saves to same location)
+        st.session_state.generated_image = "https://huggingface.co/spaces/subashdvorak/Brainstroming-FAST-API/resolve/main/brainstroming_agent/utils/image.png"
+        st.success("Image generated successfully.")
     else:
         st.error("Failed to generate image.")
-st.image("image.png",caption="Generated Image from Story")
+
+# Display image if it exists
+if st.session_state.generated_image:
+    st.image(st.session_state.generated_image, caption="Generated Image from Story")
+