Update app.py

app.py

@@ -1,59 +1,96 @@
 import streamlit as st
-from diffusers import AutoPipelineForText2Image
-import torch
+from graph import run_graph
 from PIL import Image
 import io
-import os
-import requests
-
-# --- CONFIG ---
-USE_GROQ = False  # Set to True when Groq image API is available
-GROQ_API_URL = "https://your-groq-image-api.com/generate"  # Placeholder
-
-# Force CPU
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
-@st.cache_resource
-def load_model():
-    if USE_GROQ:
-        return None  # Skip local model
-    pipe = AutoPipelineForText2Image.from_pretrained(
-        "stabilityai/sd-turbo",
-        torch_dtype=torch.float32
-    )
-    pipe.to("cpu")
-    return pipe
-
-def generate_image_local(prompt, guidance_scale):
-    pipe = load_model()
-    result = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=20)
-    return result.images[0]
-
-def generate_image_from_groq(prompt):
-    response = requests.post(GROQ_API_URL, json={"prompt": prompt})
-    if response.status_code == 200:
-        image_bytes = io.BytesIO(response.content)
-        return Image.open(image_bytes)
-    else:
-        raise Exception(f"GROQ API failed: {response.text}")
-
-# UI
-st.title("🧠 AI Image Generator (Fast with API / Groq-ready)")
-
-prompt = st.text_input("Prompt:", "A glowing alien forest with floating orbs, concept art, 8K")
-guidance = st.slider("Guidance scale", 1.0, 10.0, 3.0)
-
-if st.button("Generate"):
-    with st.spinner("Generating..."):
-        try:
-            if USE_GROQ:
-                image = generate_image_from_groq(prompt)
+
+st.set_page_config(
+    page_title="AI Image Studio",
+    layout="wide"
+)
+
+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 = []
+
+# ---------------- SIDEBAR ----------------
+st.sidebar.header("📁 Upload Files")
+
+uploaded_files = st.sidebar.file_uploader(
+    "Upload Images / Files",
+    type=["png", "jpg", "jpeg"],
+    accept_multiple_files=True
+)
+
+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:
+    st.sidebar.subheader("Preview")
+    for file in st.session_state.uploaded_files:
+        st.sidebar.image(file, caption=file.name, use_container_width=True)
+
+# ---------------- 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"])
+
+user_input = st.chat_input("Type: generate / edit / analyze image...")
+
+# ---------------- RUN AGENT ----------------
+if user_input:
+
+    st.session_state.messages.append({
+        "role": "user",
+        "content": user_input
+    })
+
+    with st.chat_message("user"):
+        st.write(user_input)
+
+    with st.chat_message("assistant"):
+        with st.spinner("AI Agent thinking..."):
+
+            result = run_graph(
+                user_input,
+                uploaded_files=st.session_state.uploaded_files
+            )
+
+            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)
+
+                buf = io.BytesIO()
+                output.save(buf, format="PNG")
+
+                st.download_button(
+                    "⬇ Download Image",
+                    data=buf.getvalue(),
+                    file_name="result.png",
+                    mime="image/png"
+                )
+
+            # ---------------- DICT RESULT (ANALYZE / OCR) ----------------
+            elif isinstance(output, dict):
+
+                st.json(output)
+
+            # ---------------- TEXT RESULT ----------------
             else:
-                image = generate_image_local(prompt, guidance)
-
-            st.image(image, caption="Generated Image", use_column_width=True)
-            buf = io.BytesIO()
-            image.save(buf, format="PNG")
-            st.download_button("Download Image", buf.getvalue(), "generated.png", "image/png")
-        except Exception as e:
-            st.error(f"Error: {e}")
+                st.write(output)
+
+    st.session_state.messages.append({
+        "role": "assistant",
+        "content": str(result.get("result"))
+    })