Update app.py

app.py

@@ -2,6 +2,8 @@ import streamlit as st
 from graph import run_graph
 from PIL import Image
 import io
+import tempfile
+import os
 
 st.set_page_config(
     page_title="AI Image Studio",
@@ -10,31 +12,55 @@ st.set_page_config(
 
 st.title("🧠 AI Image Studio (Agents + LangGraph)")
 
+
 # ---------------- SESSION STATE ----------------
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-if "uploaded_files" not in st.session_state:
-    st.session_state.uploaded_files = []
+if "uploaded_paths" not in st.session_state:
+    st.session_state.uploaded_paths = []
+
 
 # ---------------- SIDEBAR ----------------
 st.sidebar.header("📁 Upload Files")
 
 uploaded_files = st.sidebar.file_uploader(
-    "Upload Images / Files",
+    "Upload Images",
     type=["png", "jpg", "jpeg"],
     accept_multiple_files=True
 )
 
+# Convert Streamlit file → REAL FILE PATHS
 if uploaded_files:
-    st.session_state.uploaded_files = uploaded_files
-    st.sidebar.success(f"{len(uploaded_files)} file(s) uploaded")
 
-# Show uploaded images
-if st.session_state.uploaded_files:
+    paths = []
+
+    for file in uploaded_files:
+
+        temp_dir = "uploads"
+        os.makedirs(temp_dir, exist_ok=True)
+
+        file_path = os.path.join(temp_dir, file.name)
+
+        with open(file_path, "wb") as f:
+            f.write(file.getbuffer())
+
+        paths.append(file_path)
+
+    st.session_state.uploaded_paths = paths
+
+    st.sidebar.success(f"{len(paths)} files saved")
+
+
+# Preview
+if st.session_state.uploaded_paths:
+
     st.sidebar.subheader("Preview")
-    for file in st.session_state.uploaded_files:
-        st.sidebar.image(file, caption=file.name, use_container_width=True)
+
+    for path in st.session_state.uploaded_paths:
+
+        st.sidebar.image(path, caption=os.path.basename(path))
+
 
 # ---------------- CHAT UI ----------------
 st.subheader("💬 Chat with AI Agent")
@@ -43,7 +69,9 @@ for msg in st.session_state.messages:
     with st.chat_message(msg["role"]):
         st.write(msg["content"])
 
-user_input = st.chat_input("Type: generate / edit / analyze image...")
+
+user_input = st.chat_input("Generate / Edit / Analyze image...")
+
 
 # ---------------- RUN AGENT ----------------
 if user_input:
@@ -61,36 +89,41 @@ if user_input:
 
             result = run_graph(
                 user_input,
-                uploaded_files=st.session_state.uploaded_files
+                uploaded_files=st.session_state.uploaded_paths
             )
 
             output = result.get("result")
 
-            # ---------------- IMAGE RESULT ----------------
-            if hasattr(output, "show") or str(type(output)).find("PIL") != -1:
 
-                st.image(output, caption="Generated / Edited Image", use_container_width=True)
+            # ---------------- IMAGE OUTPUT ----------------
+            if isinstance(output, Image.Image):
+
+                st.image(output, caption="Result", use_container_width=True)
 
                 buf = io.BytesIO()
                 output.save(buf, format="PNG")
 
                 st.download_button(
-                    "⬇ Download Image",
-                    data=buf.getvalue(),
-                    file_name="result.png",
-                    mime="image/png"
+                    "⬇ Download",
+                    buf.getvalue(),
+                    "result.png",
+                    "image/png"
                 )
 
-            # ---------------- DICT RESULT (ANALYZE / OCR) ----------------
+
+            # ---------------- DICT OUTPUT ----------------
             elif isinstance(output, dict):
 
                 st.json(output)
 
-            # ---------------- TEXT RESULT ----------------
+
+            # ---------------- TEXT OUTPUT ----------------
             else:
+
                 st.write(output)
 
+
     st.session_state.messages.append({
         "role": "assistant",
-        "content": str(result.get("result"))
+        "content": "Done"
     })