Update app.py

app.py

@@ -7,18 +7,55 @@ from PIL import Image
 import io
 import requests
 from huggingface_hub import HfApi, login
+from pathlib import Path
+import json
 
-# Initialize session state - must be first
-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 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")
 
 class ModelGenerator:
-    @staticmethod
-    def generate_midjourney(prompt, token):
+    def __init__(self, token):
+        self.token = token
+
+    def generate_midjourney(self, prompt):
         try:
-            client = Client("mukaist/Midjourney", hf_token=token)
+            client = Client("mukaist/Midjourney", hf_token=self.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",
@@ -37,21 +74,18 @@ class ModelGenerator:
                 if isinstance(image_data, str):
                     if image_data.startswith('http'):
                         response = requests.get(image_data)
-                        image = Image.open(io.BytesIO(response.content))
+                        return ("Midjourney", Image.open(io.BytesIO(response.content)))
                     else:
-                        image = Image.open(image_data)
+                        return ("Midjourney", Image.open(image_data))
                 else:
-                    image = Image.open(io.BytesIO(image_data))
-                return ("Midjourney", image)
-            else:
-                return ("Midjourney", f"Error: Unexpected result format: {type(result)}")
+                    return ("Midjourney", Image.open(io.BytesIO(image_data)))
+            return ("Midjourney", "Error: No image generated")
         except Exception as e:
             return ("Midjourney", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_stable_cascade(prompt, token):
+    def generate_stable_cascade(self, prompt):
         try:
-            client = Client("multimodalart/stable-cascade", hf_token=token)
+            client = Client("multimodalart/stable-cascade", hf_token=self.token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -65,14 +99,23 @@ class ModelGenerator:
                 num_images_per_prompt=1,
                 api_name="/run"
             )
-            return ("Stable Cascade", result)
+            if isinstance(result, list) and len(result) > 0:
+                image_data = result[0]
+                if isinstance(image_data, str):
+                    if image_data.startswith('http'):
+                        response = requests.get(image_data)
+                        return ("Stable Cascade", Image.open(io.BytesIO(response.content)))
+                    else:
+                        return ("Stable Cascade", Image.open(image_data))
+                else:
+                    return ("Stable Cascade", Image.open(io.BytesIO(image_data)))
+            return ("Stable Cascade", "Error: No image generated")
         except Exception as e:
             return ("Stable Cascade", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_stable_diffusion_3(prompt, token):
+    def generate_stable_diffusion_3(self, prompt):
         try:
-            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=token)
+            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=self.token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -84,14 +127,15 @@ class ModelGenerator:
                 num_inference_steps=28,
                 api_name="/infer"
             )
-            return ("SD 3 Medium", result)
+            if isinstance(result, (str, bytes)):
+                return ("SD 3 Medium", Image.open(io.BytesIO(result) if isinstance(result, bytes) else result))
+            return ("SD 3 Medium", "Error: Unexpected result format")
         except Exception as e:
             return ("SD 3 Medium", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_stable_diffusion_35(prompt, token):
+    def generate_stable_diffusion_35(self, prompt):
         try:
-            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=token)
+            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=self.token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -103,15 +147,15 @@ class ModelGenerator:
                 num_inference_steps=40,
                 api_name="/infer"
             )
-            return ("SD 3.5 Large", result)
+            if isinstance(result, (str, bytes)):
+                return ("SD 3.5 Large", Image.open(io.BytesIO(result) if isinstance(result, bytes) else result))
+            return ("SD 3.5 Large", "Error: Unexpected result format")
         except Exception as e:
             return ("SD 3.5 Large", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_playground_v2_5(prompt, token):
+    def generate_playground_v2_5(self, prompt):
         try:
-            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", 
-                          hf_token=token)
+            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", hf_token=self.token)
             result = client.predict(
                 prompt,
                 prompt,  # negative prompt
@@ -123,26 +167,30 @@ class ModelGenerator:
                 True,   # randomize seed
                 api_name="/run"
             )
-            if result and isinstance(result, tuple) and result[0]:
-                return ("Playground v2.5", result[0][0]['image'])
+            if isinstance(result, tuple) and result[0] and len(result[0]) > 0:
+                image_data = result[0][0].get('image')
+                if image_data:
+                    if isinstance(image_data, str):
+                        if image_data.startswith('http'):
+                            response = requests.get(image_data)
+                            return ("Playground v2.5", Image.open(io.BytesIO(response.content)))
+                        return ("Playground v2.5", Image.open(image_data))
+                    return ("Playground v2.5", Image.open(io.BytesIO(image_data)))
             return ("Playground v2.5", "Error: No image generated")
         except Exception as e:
             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")]
-        
+def generate_images(prompt, selected_models, token):
     results = []
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = []
+        generator = ModelGenerator(token)
         model_map = {
-            "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)
+            "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
         }
         
         for model in selected_models:
@@ -150,17 +198,20 @@ def generate_images(prompt, selected_models):
                 futures.append(executor.submit(model_map[model], prompt))
         
         for future in concurrent.futures.as_completed(futures):
-            results.append(future.result())
+            try:
+                result = future.result()
+                if result:
+                    results.append(result)
+            except Exception as e:
+                st.error(f"Error during image generation: {str(e)}")
     
     return results
 
 def handle_prompt_click(prompt_text, key):
-    if not st.session_state.get('is_authenticated') or not st.session_state.get('hf_token'):
+    if not st.session_state.is_authenticated:
         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:
@@ -168,42 +219,31 @@ def handle_prompt_click(prompt_text, key):
         return
 
     with st.spinner('Generating artwork...'):
-        results = generate_images(prompt_text, selected_models)
-        st.session_state[f'generated_images_{key}'] = results
-        st.success("Artwork generated successfully!")
+        results = generate_images(prompt_text, selected_models, st.session_state.hf_token)
+        if results:
+            st.session_state[f'generated_images_{key}'] = results
+            st.success("Artwork generated successfully!")
+            
+            # Display images immediately
+            cols = st.columns(len(results))
+            for col, (model_name, result) in zip(cols, results):
+                with col:
+                    st.markdown(f"**{model_name}**")
+                    if isinstance(result, str) and result.startswith("Error"):
+                        st.error(result)
+                    elif isinstance(result, Image.Image):
+                        st.image(result, use_container_width=True)
+                    else:
+                        st.error(f"Unexpected result type: {type(result)}")
 
 def main():
     st.title("🎨 Multi-Model Art Generator")
+    
+    init_session_state()
+    authenticate_user()
 
-    # Handle authentication in sidebar
-    with st.sidebar:
-        st.header("🔐 Authentication")
-        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.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.get('is_authenticated') and st.session_state.get('hf_token'):
-            st.markdown("---")
+    if st.session_state.is_authenticated:
+        with st.sidebar:
             st.header("Model Selection")
             st.session_state['selected_models'] = st.multiselect(
                 "Choose AI Models",
@@ -226,8 +266,6 @@ def main():
             - **Playground v2.5**: Advanced model with high customization
             """)
 
-    # Only show the main interface if 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 = {
@@ -278,23 +316,25 @@ def main():
         st.markdown("---")
         st.markdown("### Generated Artwork:")
         
+        # Display any previously generated images
         for key in st.session_state:
-            if key.startswith('selected_prompt_'):
+            if key.startswith('generated_images_'):
                 idx = key.split('_')[-1]
-                images_key = f'generated_images_{idx}'
+                prompt_key = f'selected_prompt_{idx}'
                 
-                if images_key in st.session_state:
-                    st.write("Prompt:", st.session_state[key])
+                if prompt_key in st.session_state:
+                    st.write("Prompt:", st.session_state[prompt_key])
                     
-                    cols = st.columns(len(st.session_state[images_key]))
-                    
-                    for col, (model_name, result) in zip(cols, st.session_state[images_key]):
+                    cols = st.columns(len(st.session_state[key]))
+                    for col, (model_name, result) in zip(cols, st.session_state[key]):
                         with col:
                             st.markdown(f"**{model_name}**")
                             if isinstance(result, str) and result.startswith("Error"):
                                 st.error(result)
-                            else:
+                            elif isinstance(result, Image.Image):
                                 st.image(result, use_container_width=True)
+                            else:
+                                st.error(f"Unexpected result type: {type(result)}")
     else:
         st.info("Please login with your HuggingFace account to use the app")