Updated11

app.py

@@ -104,32 +104,32 @@ def generate_image(pipe, prompt, num_inference_steps=30):
 def main():
     st.title("🎨 Stable Diffusion Image Generator")
     
+    # Initialize session state
+    if 'pipe' not in st.session_state:
+        st.session_state.pipe = None
+    if 'model_initialized' not in st.session_state:
+        st.session_state.model_initialized = False
+    
     # Sidebar for configuration
     st.sidebar.header("Configuration")
     
     # HuggingFace token input
     hf_token = st.sidebar.text_input("Enter HuggingFace Token:", type="password")
     
-    # Initialize session state
-    if 'pipe' not in st.session_state:
-        st.session_state.pipe = None
-    
-    # Initialize button
-    if st.sidebar.button("Initialize Model"):
-        if not hf_token:
-            st.sidebar.error("Please enter your HuggingFace token")
-        else:
-            st.session_state.pipe = setup_stable_diffusion(hf_token)
-            if st.session_state.pipe is not None:
-                st.sidebar.success("Model initialized successfully!")
+    # Auto-initialize when token is entered
+    if hf_token and not st.session_state.model_initialized:
+        st.session_state.pipe = setup_stable_diffusion(hf_token)
+        if st.session_state.pipe is not None:
+            st.session_state.model_initialized = True
+            st.sidebar.success("Model initialized successfully!")
     
     # Main interface
     prompt = st.text_input("Enter your prompt:", "A stunning landscape with mountains and a lake at sunset")
     num_steps = st.slider("Number of inference steps:", min_value=20, max_value=100, value=30)
     
     if st.button("Generate Image"):
-        if st.session_state.pipe is None:
-            st.error("Please initialize the model first using the sidebar!")
+        if not st.session_state.model_initialized:
+            st.error("Please enter your HuggingFace token in the sidebar to initialize the model!")
         else:
             try:
                 image = generate_image(st.session_state.pipe, prompt, num_steps)
@@ -151,4 +151,5 @@ def main():
                     torch.cuda.empty_cache()
 
 if __name__ == "__main__":
-    main()
+    main()
+