Update app.py

app.py

@@ -16,12 +16,9 @@ if 'is_authenticated' not in st.session_state:
 
 class ModelGenerator:
     @staticmethod
-    def generate_midjourney(prompt):
-        if not st.session_state.get('hf_token'):
-            return ("Midjourney", "Error: No authentication token found")
-            
+    def generate_midjourney(prompt, token):
         try:
-            client = Client("mukaist/Midjourney", hf_token=st.session_state['hf_token'])
+            client = Client("mukaist/Midjourney", hf_token=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",
@@ -52,12 +49,9 @@ class ModelGenerator:
             return ("Midjourney", f"Error: {str(e)}")
 
     @staticmethod
-    def generate_stable_cascade(prompt):
-        if not st.session_state.get('hf_token'):
-            return ("Stable Cascade", "Error: No authentication token found")
-            
+    def generate_stable_cascade(prompt, token):
         try:
-            client = Client("multimodalart/stable-cascade", hf_token=st.session_state['hf_token'])
+            client = Client("multimodalart/stable-cascade", hf_token=token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -76,12 +70,9 @@ class ModelGenerator:
             return ("Stable Cascade", f"Error: {str(e)}")
 
     @staticmethod
-    def generate_stable_diffusion_3(prompt):
-        if not st.session_state.get('hf_token'):
-            return ("SD 3 Medium", "Error: No authentication token found")
-            
+    def generate_stable_diffusion_3(prompt, token):
         try:
-            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=st.session_state['hf_token'])
+            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -98,12 +89,9 @@ class ModelGenerator:
             return ("SD 3 Medium", f"Error: {str(e)}")
 
     @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")
-            
+    def generate_stable_diffusion_35(prompt, token):
         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=token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -120,13 +108,10 @@ class ModelGenerator:
             return ("SD 3.5 Large", f"Error: {str(e)}")
 
     @staticmethod
-    def generate_playground_v2_5(prompt):
-        if not st.session_state.get('hf_token'):
-            return ("Playground v2.5", "Error: No authentication token found")
-            
+    def generate_playground_v2_5(prompt, token):
         try:
             client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", 
-                          hf_token=st.session_state['hf_token'])
+                          hf_token=token)
             result = client.predict(
                 prompt,
                 prompt,  # negative prompt
@@ -145,15 +130,19 @@ class ModelGenerator:
             return ("Playground v2.5", f"Error: {str(e)}")
 
 def generate_images(prompt, selected_models):
+    token = st.session_state.get('hf_token')
+    if not token:
+        return [("Error", "No authentication token found")]
+        
     results = []
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = []
         model_map = {
-            "Midjourney": ModelGenerator.generate_midjourney,
-            "Stable Cascade": ModelGenerator.generate_stable_cascade,
-            "SD 3 Medium": ModelGenerator.generate_stable_diffusion_3,
-            "SD 3.5 Large": ModelGenerator.generate_stable_diffusion_35,
-            "Playground v2.5": ModelGenerator.generate_playground_v2_5
+            "Midjourney": lambda p: ModelGenerator.generate_midjourney(p, token),
+            "Stable Cascade": lambda p: ModelGenerator.generate_stable_cascade(p, token),
+            "SD 3 Medium": lambda p: ModelGenerator.generate_stable_diffusion_3(p, token),
+            "SD 3.5 Large": lambda p: ModelGenerator.generate_stable_diffusion_35(p, token),
+            "Playground v2.5": lambda p: ModelGenerator.generate_playground_v2_5(p, token)
         }
         
         for model in selected_models:
@@ -204,7 +193,6 @@ def main():
                         # Verify token is valid
                         api = HfApi(token=token)
                         api.whoami()
-                        # 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!")