Update app.py

app.py

@@ -435,12 +435,19 @@ def generate_diagram(prompt: str, width: int, height: int, seed: int = None) ->
 
 
 def main():
-    st.set_page_config(page_title="FLUX Diagram Generator", layout="wide")
+    st.set_page_config(page_title="FLUX Diagram Generator", layout="wide", initial_sidebar_state="expanded")
     
+    if "initialized" not in st.session_state:
+        st.session_state.initialized = True
+        st.session_state.generated_images = {}
+
     st.title("🎨 FLUX Diagram Generator")
     st.markdown("Generate beautiful hand-drawn style diagrams using FLUX AI")
 
-    # Sidebar for examples
+    # Initialize client
+    client = Client("black-forest-labs/FLUX.1-schnell", hf_token=HF_TOKEN)
+
+    # Sidebar with examples
     st.sidebar.title("📚 Example Templates")
     selected_example = st.sidebar.selectbox(
         "Choose a template",
@@ -448,28 +455,33 @@ def main():
         format_func=lambda x: DIAGRAM_EXAMPLES[x]["title"]
     )
 
-    # Main content area
+    # Main content
     col1, col2 = st.columns([2, 1])
 
     with col1:
-        # Input area
-        prompt = st.text_area(
-            "Diagram Prompt",
-            value=DIAGRAM_EXAMPLES[selected_example]["prompt"],
-            height=200
-        )
-
-        # Configuration
+        example = DIAGRAM_EXAMPLES[selected_example]
+        prompt = st.text_area("Diagram Prompt", value=example["prompt"], height=200)
+        
         with st.expander("Advanced Configuration"):
             width = st.number_input("Width", min_value=512, max_value=2048, value=1024, step=128)
             height = st.number_input("Height", min_value=512, max_value=2048, value=1024, step=128)
-            seed = st.number_input("Seed (optional)", value=None, step=1)
 
         if st.button("🎨 Generate Diagram"):
-            with st.spinner("Generating your diagram..."):
-                result = generate_diagram(prompt, width, height, seed)
-                if result:
+            try:
+                with st.spinner("Generating diagram..."):
+                    result = client.predict(
+                        prompt,
+                        1872187377,  # seed
+                        False,       # randomize_seed
+                        width,
+                        height,
+                        4,          # num_inference_steps
+                        api_name="/infer"
+                    )
+                    st.session_state.generated_images[prompt] = result
                     st.image(result, caption="Generated Diagram", use_column_width=True)
+            except Exception as e:
+                st.error(f"Error generating diagram: {str(e)}")
 
     with col2:
         st.subheader("Tips for Better Results")