Update app.py

app.py

@@ -8,17 +8,20 @@ import io
 import requests
 from huggingface_hub import HfApi, login
 
-# Initialize session state for authentication
+# Initialize session state - must be first
 if 'hf_token' not in st.session_state:
-    st.session_state.hf_token = None
+    st.session_state['hf_token'] = None
 if 'is_authenticated' not in st.session_state:
-    st.session_state.is_authenticated = False
+    st.session_state['is_authenticated'] = False
 
 class ModelGenerator:
     @staticmethod
     def generate_midjourney(prompt):
+        if not st.session_state.get('hf_token'):
+            return ("Midjourney", "Error: No authentication token found")
+            
         try:
-            client = Client("mukaist/Midjourney", hf_token=st.session_state.hf_token)
+            client = Client("mukaist/Midjourney", hf_token=st.session_state['hf_token'])
             result = client.predict(
                 prompt=prompt,
                 negative_prompt="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
@@ -48,11 +51,13 @@ class ModelGenerator:
         except Exception as e:
             return ("Midjourney", f"Error: {str(e)}")
 
-    # Other model generators remain exactly the same, just changing HF_TOKEN to st.session_state.hf_token
     @staticmethod
     def generate_stable_cascade(prompt):
+        if not st.session_state.get('hf_token'):
+            return ("Stable Cascade", "Error: No authentication token found")
+            
         try:
-            client = Client("multimodalart/stable-cascade", hf_token=st.session_state.hf_token)
+            client = Client("multimodalart/stable-cascade", hf_token=st.session_state['hf_token'])
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -72,8 +77,11 @@ class ModelGenerator:
 
     @staticmethod
     def generate_stable_diffusion_3(prompt):
+        if not st.session_state.get('hf_token'):
+            return ("SD 3 Medium", "Error: No authentication token found")
+            
         try:
-            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=st.session_state.hf_token)
+            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=st.session_state['hf_token'])
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -91,8 +99,11 @@ class ModelGenerator:
 
     @staticmethod
     def generate_stable_diffusion_35(prompt):
+        if not st.session_state.get('hf_token'):
+            return ("SD 3.5 Large", "Error: No authentication token found")
+            
         try:
-            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=st.session_state.hf_token)
+            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=st.session_state['hf_token'])
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -110,9 +121,12 @@ class ModelGenerator:
 
     @staticmethod
     def generate_playground_v2_5(prompt):
+        if not st.session_state.get('hf_token'):
+            return ("Playground v2.5", "Error: No authentication token found")
+            
         try:
             client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", 
-                          hf_token=st.session_state.hf_token)
+                          hf_token=st.session_state['hf_token'])
             result = client.predict(
                 prompt,
                 prompt,  # negative prompt
@@ -152,7 +166,7 @@ def generate_images(prompt, selected_models):
     return results
 
 def handle_prompt_click(prompt_text, key):
-    if not st.session_state.is_authenticated:
+    if not st.session_state.get('is_authenticated') or not st.session_state.get('hf_token'):
         st.error("Please login with your HuggingFace account first!")
         return
         
@@ -175,11 +189,11 @@ def main():
     # Handle authentication in sidebar
     with st.sidebar:
         st.header("🔐 Authentication")
-        if st.session_state.is_authenticated:
+        if st.session_state.get('is_authenticated') and st.session_state.get('hf_token'):
             st.success("✓ Logged in to HuggingFace")
             if st.button("Logout"):
-                st.session_state.hf_token = None
-                st.session_state.is_authenticated = False
+                st.session_state['hf_token'] = None
+                st.session_state['is_authenticated'] = False
                 st.rerun()
         else:
             token = st.text_input("Enter HuggingFace Token", type="password",
@@ -190,8 +204,9 @@ def main():
                         # Verify token is valid
                         api = HfApi(token=token)
                         api.whoami()
-                        st.session_state.hf_token = token
-                        st.session_state.is_authenticated = True
+                        # Use dictionary style access to ensure proper setting
+                        st.session_state['hf_token'] = token
+                        st.session_state['is_authenticated'] = True
                         st.success("Successfully logged in!")
                         st.rerun()
                     except Exception as e:
@@ -199,7 +214,7 @@ def main():
                 else:
                     st.error("Please enter your HuggingFace token")
         
-        if st.session_state.is_authenticated:
+        if st.session_state.get('is_authenticated') and st.session_state.get('hf_token'):
             st.markdown("---")
             st.header("Model Selection")
             st.session_state['selected_models'] = st.multiselect(
@@ -224,7 +239,7 @@ def main():
             """)
 
     # Only show the main interface if authenticated
-    if st.session_state.is_authenticated:
+    if st.session_state.get('is_authenticated') and st.session_state.get('hf_token'):
         st.markdown("### Select a prompt style to generate artwork:")
 
         prompt_emojis = {