Update app.py

app.py

@@ -1,9 +1,13 @@
+import os
+import io
+from PIL import Image
 import streamlit as st
+
 from graph import run_graph
-from PIL import Image
-import io
-import tempfile
-import os
+
+# -------------------------------------------------------
+# PAGE CONFIG
+# -------------------------------------------------------
 
 st.set_page_config(
     page_title="AI Image Studio",
@@ -12,17 +16,27 @@ st.set_page_config(
 
 st.title("🧠 AI Image Studio (Agents + LangGraph)")
 
+# -------------------------------------------------------
+# SESSION STATE
+# -------------------------------------------------------
 
-# ---------------- SESSION STATE ----------------
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
 if "uploaded_paths" not in st.session_state:
     st.session_state.uploaded_paths = []
 
+if "analysis" not in st.session_state:
+    st.session_state.analysis = None
 
-# ---------------- SIDEBAR ----------------
-st.sidebar.header("📁 Upload Files")
+if "last_uploaded" not in st.session_state:
+    st.session_state.last_uploaded = []
+
+# -------------------------------------------------------
+# SIDEBAR
+# -------------------------------------------------------
+
+st.sidebar.header("📁 Upload Images")
 
 uploaded_files = st.sidebar.file_uploader(
     "Upload Images",
@@ -30,62 +44,176 @@ uploaded_files = st.sidebar.file_uploader(
     accept_multiple_files=True
 )
 
-# Convert Streamlit file → REAL FILE PATHS
+# -------------------------------------------------------
+# SAVE FILES
+# -------------------------------------------------------
+
 if uploaded_files:
 
+    os.makedirs("uploads", exist_ok=True)
+
     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)
+        file_path = os.path.join("uploads", file.name)
 
         with open(file_path, "wb") as f:
             f.write(file.getbuffer())
 
         paths.append(file_path)
 
+    # Save paths
     st.session_state.uploaded_paths = paths
 
-    st.sidebar.success(f"{len(paths)} files saved")
+    # ---------------------------------------------------
+    # AUTO ANALYZE ONLY IF NEW IMAGE
+    # ---------------------------------------------------
+
+    if paths != st.session_state.last_uploaded:
+
+        st.session_state.last_uploaded = paths
+
+        with st.spinner("🔍 Analyzing image..."):
+
+            result = run_graph(
+                "analyze this image",
+                uploaded_files=paths
+            )
+
+        analysis = result.get("result")
+
+        st.session_state.analysis = analysis
+
+        if isinstance(analysis, dict):
+
+            description = analysis.get(
+                "description",
+                "Image analyzed."
+            )
+
+            resolution = analysis.get("resolution", "")
+
+            style = analysis.get("style", "")
+
+            lighting = analysis.get("lighting", "")
+
+            objects = analysis.get("objects", [])
+
+            message = f"""📷 **Image Analysis**
+
+**Description**
+
+{description}
+"""
 
+            if resolution:
+                message += f"\n**Resolution:** {resolution}"
+
+            if style:
+                message += f"\n\n**Style:** {style}"
+
+            if lighting:
+                message += f"\n\n**Lighting:** {lighting}"
+
+            if objects:
+                message += f"\n\n**Objects:** {', '.join(objects)}"
+
+            message += """
+
+---
+
+### What would you like to do?
+
+• Remove background
+
+• Make it Anime
+
+• Make it Pixar style
+
+• Enhance quality
+
+• Replace objects
+
+• Change colors
+
+• Add new objects
+
+• Extract text
+"""
+
+            st.session_state.messages.append(
+                {
+                    "role": "assistant",
+                    "content": message
+                }
+            )
+
+        else:
+
+            st.session_state.messages.append(
+                {
+                    "role": "assistant",
+                    "content": str(analysis)
+                }
+            )
+
+# -------------------------------------------------------
+# SIDEBAR PREVIEW
+# -------------------------------------------------------
 
-# Preview
 if st.session_state.uploaded_paths:
 
-    st.sidebar.subheader("Preview")
+    st.sidebar.subheader("🖼 Preview")
 
-    for path in st.session_state.uploaded_paths:
+    for img in st.session_state.uploaded_paths:
 
-        st.sidebar.image(path, caption=os.path.basename(path))
+        st.sidebar.image(
+            img,
+            caption=os.path.basename(img),
+            use_container_width=True
+        )
 
+# -------------------------------------------------------
+# CHAT
+# -------------------------------------------------------
 
-# ---------------- CHAT UI ----------------
 st.subheader("💬 Chat with AI Agent")
 
 for msg in st.session_state.messages:
+
     with st.chat_message(msg["role"]):
-        st.write(msg["content"])
 
+        st.markdown(msg["content"])
+
+# -------------------------------------------------------
+# USER INPUT
+# -------------------------------------------------------
 
-user_input = st.chat_input("Generate / Edit / Analyze image...")
+user_input = st.chat_input(
+    "Generate / Edit / Analyze image..."
+)
 
+# -------------------------------------------------------
+# RUN AGENT
+# -------------------------------------------------------
 
-# ---------------- RUN AGENT ----------------
 if user_input:
 
-    st.session_state.messages.append({
-        "role": "user",
-        "content": user_input
-    })
+    st.session_state.messages.append(
+        {
+            "role": "user",
+            "content": user_input
+        }
+    )
 
     with st.chat_message("user"):
-        st.write(user_input)
+
+        st.markdown(user_input)
 
     with st.chat_message("assistant"):
-        with st.spinner("AI Agent thinking..."):
+
+        with st.spinner("🤖 AI Agent thinking..."):
 
             result = run_graph(
                 user_input,
@@ -94,36 +222,59 @@ if user_input:
 
             output = result.get("result")
 
+            assistant_message = ""
+
+            # ------------------------------------------------
+            # IMAGE OUTPUT
+            # ------------------------------------------------
 
-            # ---------------- IMAGE OUTPUT ----------------
             if isinstance(output, Image.Image):
 
-                st.image(output, caption="Result", use_container_width=True)
+                st.image(
+                    output,
+                    caption="Generated Image",
+                    use_container_width=True
+                )
+
+                buffer = io.BytesIO()
 
-                buf = io.BytesIO()
-                output.save(buf, format="PNG")
+                output.save(buffer, format="PNG")
 
                 st.download_button(
-                    "⬇ Download",
-                    buf.getvalue(),
+                    "⬇ Download Image",
+                    buffer.getvalue(),
                     "result.png",
                     "image/png"
                 )
 
+                assistant_message = "✅ Image generated successfully."
+
+            # ------------------------------------------------
+            # DICTIONARY OUTPUT
+            # ------------------------------------------------
 
-            # ---------------- DICT OUTPUT ----------------
             elif isinstance(output, dict):
 
                 st.json(output)
 
+                assistant_message = output.get(
+                    "description",
+                    str(output)
+                )
+
+            # ------------------------------------------------
+            # TEXT OUTPUT
+            # ------------------------------------------------
 
-            # ---------------- TEXT OUTPUT ----------------
             else:
 
-                st.write(output)
+                st.markdown(str(output))
 
+                assistant_message = str(output)
 
-    st.session_state.messages.append({
-        "role": "assistant",
-        "content": "Done"
-    })
+    st.session_state.messages.append(
+        {
+            "role": "assistant",
+            "content": assistant_message
+        }
+    )