Update app.py

app.py

@@ -7,55 +7,18 @@ from PIL import Image
 import io
 import requests
 from huggingface_hub import HfApi, login
-from pathlib import Path
-import json
 
-def init_session_state():
-    """Initialize session state variables"""
-    if 'hf_token' not in st.session_state:
-        st.session_state.hf_token = None
-    if 'is_authenticated' not in st.session_state:
-        st.session_state.is_authenticated = False
-
-def save_token(token):
-    """Save token to session state"""
-    st.session_state.hf_token = token
-    st.session_state.is_authenticated = True
-
-def authenticate_user():
-    """Handle user authentication with HuggingFace"""
-    st.sidebar.markdown("## 🔐 Authentication")
-    
-    if st.session_state.is_authenticated:
-        st.sidebar.success("✓ Logged in to HuggingFace")
-        if st.sidebar.button("Logout"):
-            st.session_state.hf_token = None
-            st.session_state.is_authenticated = False
-            st.rerun()
-    else:
-        token = st.sidebar.text_input("Enter HuggingFace Token", type="password", 
-                                    help="Get your token from https://huggingface.co/settings/tokens")
-        if st.sidebar.button("Login"):
-            if token:
-                try:
-                    # Verify token is valid
-                    api = HfApi(token=token)
-                    api.whoami()
-                    save_token(token)
-                    st.sidebar.success("Successfully logged in!")
-                    st.rerun()
-                except Exception as e:
-                    st.sidebar.error(f"Authentication failed: {str(e)}")
-            else:
-                st.sidebar.error("Please enter your HuggingFace token")
+# Initialize session state for authentication
+if 'hf_token' not in st.session_state:
+    st.session_state.hf_token = None
+if 'is_authenticated' not in st.session_state:
+    st.session_state.is_authenticated = False
 
 class ModelGenerator:
-    def __init__(self, token):
-        self.token = token
-
-    def generate_midjourney(self, prompt):
+    @staticmethod
+    def generate_midjourney(prompt):
         try:
-            client = Client("mukaist/Midjourney", hf_token=self.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",
@@ -85,9 +48,11 @@ class ModelGenerator:
         except Exception as e:
             return ("Midjourney", f"Error: {str(e)}")
 
-    def generate_stable_cascade(self, prompt):
+    # Other model generators remain exactly the same, just changing HF_TOKEN to st.session_state.hf_token
+    @staticmethod
+    def generate_stable_cascade(prompt):
         try:
-            client = Client("multimodalart/stable-cascade", hf_token=self.token)
+            client = Client("multimodalart/stable-cascade", hf_token=st.session_state.hf_token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -105,9 +70,10 @@ class ModelGenerator:
         except Exception as e:
             return ("Stable Cascade", f"Error: {str(e)}")
 
-    def generate_stable_diffusion_3(self, prompt):
+    @staticmethod
+    def generate_stable_diffusion_3(prompt):
         try:
-            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=self.token)
+            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=st.session_state.hf_token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -123,9 +89,10 @@ class ModelGenerator:
         except Exception as e:
             return ("SD 3 Medium", f"Error: {str(e)}")
 
-    def generate_stable_diffusion_35(self, prompt):
+    @staticmethod
+    def generate_stable_diffusion_35(prompt):
         try:
-            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=self.token)
+            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=st.session_state.hf_token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -141,9 +108,11 @@ class ModelGenerator:
         except Exception as e:
             return ("SD 3.5 Large", f"Error: {str(e)}")
 
-    def generate_playground_v2_5(self, prompt):
+    @staticmethod
+    def generate_playground_v2_5(prompt):
         try:
-            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", hf_token=self.token)
+            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", 
+                          hf_token=st.session_state.hf_token)
             result = client.predict(
                 prompt,
                 prompt,  # negative prompt
@@ -161,17 +130,16 @@ class ModelGenerator:
         except Exception as e:
             return ("Playground v2.5", f"Error: {str(e)}")
 
-def generate_images(prompt, selected_models, token):
+def generate_images(prompt, selected_models):
     results = []
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = []
-        generator = ModelGenerator(token)
         model_map = {
-            "Midjourney": generator.generate_midjourney,
-            "Stable Cascade": generator.generate_stable_cascade,
-            "SD 3 Medium": generator.generate_stable_diffusion_3,
-            "SD 3.5 Large": generator.generate_stable_diffusion_35,
-            "Playground v2.5": generator.generate_playground_v2_5
+            "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
         }
         
         for model in selected_models:
@@ -188,6 +156,8 @@ def handle_prompt_click(prompt_text, key):
         st.error("Please login with your HuggingFace account first!")
         return
         
+    st.session_state[f'selected_prompt_{key}'] = prompt_text
+    
     selected_models = st.session_state.get('selected_models', [])
     
     if not selected_models:
@@ -195,22 +165,42 @@ def handle_prompt_click(prompt_text, key):
         return
 
     with st.spinner('Generating artwork...'):
-        results = generate_images(prompt_text, selected_models, st.session_state.hf_token)
+        results = generate_images(prompt_text, selected_models)
         st.session_state[f'generated_images_{key}'] = results
         st.success("Artwork generated successfully!")
 
 def main():
     st.title("🎨 Multi-Model Art Generator")
-    
-    # Initialize session state
-    init_session_state()
-    
-    # Handle authentication
-    authenticate_user()
 
-    # Only show the main interface if authenticated
-    if st.session_state.is_authenticated:
-        with st.sidebar:
+    # Handle authentication in sidebar
+    with st.sidebar:
+        st.header("🔐 Authentication")
+        if st.session_state.is_authenticated:
+            st.success("✓ Logged in to HuggingFace")
+            if st.button("Logout"):
+                st.session_state.hf_token = None
+                st.session_state.is_authenticated = False
+                st.rerun()
+        else:
+            token = st.text_input("Enter HuggingFace Token", type="password",
+                                help="Get your token from https://huggingface.co/settings/tokens")
+            if st.button("Login"):
+                if token:
+                    try:
+                        # Verify token is valid
+                        api = HfApi(token=token)
+                        api.whoami()
+                        st.session_state.hf_token = token
+                        st.session_state.is_authenticated = True
+                        st.success("Successfully logged in!")
+                        st.rerun()
+                    except Exception as e:
+                        st.error(f"Authentication failed: {str(e)}")
+                else:
+                    st.error("Please enter your HuggingFace token")
+        
+        if st.session_state.is_authenticated:
+            st.markdown("---")
             st.header("Model Selection")
             st.session_state['selected_models'] = st.multiselect(
                 "Choose AI Models",
@@ -233,6 +223,8 @@ def main():
             - **Playground v2.5**: Advanced model with high customization
             """)
 
+    # Only show the main interface if authenticated
+    if st.session_state.is_authenticated:
         st.markdown("### Select a prompt style to generate artwork:")
 
         prompt_emojis = {