Update app.py

app.py

@@ -6,17 +6,56 @@ import os
 from PIL import Image
 import io
 import requests
+from huggingface_hub import HfApi, login
+from pathlib import Path
+import json
 
-# Get token from environment variable
-HF_TOKEN = os.getenv('ArtToken')
-if not HF_TOKEN:
-    raise ValueError("Please set the 'ArtToken' environment variable with your Hugging Face token")
+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.experimental_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.experimental_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):
+    def __init__(self, token):
+        self.token = token
+
+    def generate_midjourney(self, prompt):
         try:
-            client = Client("mukaist/Midjourney", hf_token=HF_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",
@@ -30,20 +69,15 @@ class ModelGenerator:
                 api_name="/run"
             )
             
-            # Handle the result based on its type
             if isinstance(result, list) and len(result) > 0:
-                # If result is a list of file paths or URLs
                 image_data = result[0]
                 if isinstance(image_data, str):
                     if image_data.startswith('http'):
-                        # If it's a URL, download the image
                         response = requests.get(image_data)
                         image = Image.open(io.BytesIO(response.content))
                     else:
-                        # If it's a file path
                         image = Image.open(image_data)
                 else:
-                    # If it's already image data
                     image = Image.open(io.BytesIO(image_data))
                 return ("Midjourney", image)
             else:
@@ -51,10 +85,9 @@ class ModelGenerator:
         except Exception as e:
             return ("Midjourney", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_stable_cascade(prompt):
+    def generate_stable_cascade(self, prompt):
         try:
-            client = Client("multimodalart/stable-cascade", hf_token=HF_TOKEN)
+            client = Client("multimodalart/stable-cascade", hf_token=self.token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -72,10 +105,9 @@ class ModelGenerator:
         except Exception as e:
             return ("Stable Cascade", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_stable_diffusion_3(prompt):
+    def generate_stable_diffusion_3(self, prompt):
         try:
-            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=HF_TOKEN)
+            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=self.token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -91,10 +123,9 @@ class ModelGenerator:
         except Exception as e:
             return ("SD 3 Medium", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_stable_diffusion_35(prompt):
+    def generate_stable_diffusion_35(self, prompt):
         try:
-            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=HF_TOKEN)
+            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=self.token)
             result = client.predict(
                 prompt=prompt,
                 negative_prompt=prompt,
@@ -110,10 +141,9 @@ class ModelGenerator:
         except Exception as e:
             return ("SD 3.5 Large", f"Error: {str(e)}")
 
-    @staticmethod
-    def generate_playground_v2_5(prompt):
+    def generate_playground_v2_5(self, prompt):
         try:
-            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", hf_token=HF_TOKEN)
+            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", hf_token=self.token)
             result = client.predict(
                 prompt,
                 prompt,  # negative prompt
@@ -125,23 +155,23 @@ class ModelGenerator:
                 True,   # randomize seed
                 api_name="/run"
             )
-            # Result is a tuple (gallery, seed), we want just the first image from gallery
             if result and isinstance(result, tuple) and result[0]:
                 return ("Playground v2.5", result[0][0]['image'])
             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):
+def generate_images(prompt, selected_models, token):
     results = []
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = []
+        generator = ModelGenerator(token)
         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": 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:
@@ -154,12 +184,10 @@ def generate_images(prompt, selected_models):
     return results
 
 def handle_prompt_click(prompt_text, key):
-    if not HF_TOKEN:
-        st.error("Environment variable 'ArtToken' is not set!")
+    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:
@@ -167,113 +195,113 @@ def handle_prompt_click(prompt_text, key):
         return
 
     with st.spinner('Generating artwork...'):
-        results = generate_images(prompt_text, selected_models)
+        results = generate_images(prompt_text, selected_models, st.session_state.hf_token)
         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()
 
-    with st.sidebar:
-        st.header("Configuration")
-        
-        # Show token status
-        if HF_TOKEN:
-            st.success("✓ ArtToken loaded from environment")
-        else:
-            st.error("⚠ ArtToken not found in environment")
-        
-        st.markdown("---")
-        st.header("Model Selection")
-        st.session_state['selected_models'] = st.multiselect(
-            "Choose AI Models",
-            ["Midjourney", "Stable Cascade", "SD 3 Medium", "SD 3.5 Large", "Playground v2.5"],
-            default=["Midjourney"]
-        )
-        
-        st.markdown("---")
-        st.markdown("### Selected Models:")
-        for model in st.session_state['selected_models']:
-            st.write(f"✓ {model}")
-        
-        st.markdown("---")
-        st.markdown("### Model Information:")
-        st.markdown("""
-        - **Midjourney**: Best for artistic and creative imagery
-        - **Stable Cascade**: New architecture with high detail
-        - **SD 3 Medium**: Fast and efficient generation
-        - **SD 3.5 Large**: Highest quality, slower generation
-        - **Playground v2.5**: Advanced model with high customization
-        """)
+    # Only show the main interface if authenticated
+    if st.session_state.is_authenticated:
+        with st.sidebar:
+            st.header("Model Selection")
+            st.session_state['selected_models'] = st.multiselect(
+                "Choose AI Models",
+                ["Midjourney", "Stable Cascade", "SD 3 Medium", "SD 3.5 Large", "Playground v2.5"],
+                default=["Midjourney"]
+            )
+            
+            st.markdown("---")
+            st.markdown("### Selected Models:")
+            for model in st.session_state['selected_models']:
+                st.write(f"✓ {model}")
+            
+            st.markdown("---")
+            st.markdown("### Model Information:")
+            st.markdown("""
+            - **Midjourney**: Best for artistic and creative imagery
+            - **Stable Cascade**: New architecture with high detail
+            - **SD 3 Medium**: Fast and efficient generation
+            - **SD 3.5 Large**: Highest quality, slower generation
+            - **Playground v2.5**: Advanced model with high customization
+            """)
 
-    st.markdown("### Select a prompt style to generate artwork:")
+        st.markdown("### Select a prompt style to generate artwork:")
 
-    prompt_emojis = {
-        "AIart/AIArtistCommunity": "🤖",
-        "Black & White": "⚫⚪",
-        "Black & Yellow": "⚫💛",
-        "Blindfold": "🙈",
-        "Break": "💔",
-        "Broken": "🔨",
-        "Christmas Celebrations art": "🎄",
-        "Colorful Art": "🎨",
-        "Crimson art": "🔴",
-        "Eyes Art": "👁️",
-        "Going out with Style": "💃",
-        "Hooded Girl": "🧥",
-        "Lips": "👄",
-        "MAEKHLONG": "🏮",
-        "Mermaid": "🧜‍♀️",
-        "Morning Sunshine": "🌅",
-        "Music Art": "🎵",
-        "Owl": "🦉",
-        "Pink": "💗",
-        "Purple": "💜",
-        "Rain": "🌧️",
-        "Red Moon": "🌑",
-        "Rose": "🌹",
-        "Snow": "❄️",
-        "Spacesuit Girl": "👩‍🚀",
-        "Steampunk": "⚙️",
-        "Succubus": "😈",
-        "Sunlight": "☀️",
-        "Weird art": "🎭",
-        "White Hair": "👱‍♀️",
-        "Wings art": "👼",
-        "Woman with Sword": "⚔️"
-    }
+        prompt_emojis = {
+            "AIart/AIArtistCommunity": "🤖",
+            "Black & White": "⚫⚪",
+            "Black & Yellow": "⚫💛",
+            "Blindfold": "🙈",
+            "Break": "💔",
+            "Broken": "🔨",
+            "Christmas Celebrations art": "🎄",
+            "Colorful Art": "🎨",
+            "Crimson art": "🔴",
+            "Eyes Art": "👁️",
+            "Going out with Style": "💃",
+            "Hooded Girl": "🧥",
+            "Lips": "👄",
+            "MAEKHLONG": "🏮",
+            "Mermaid": "🧜‍♀️",
+            "Morning Sunshine": "🌅",
+            "Music Art": "🎵",
+            "Owl": "🦉",
+            "Pink": "💗",
+            "Purple": "💜",
+            "Rain": "🌧️",
+            "Red Moon": "🌑",
+            "Rose": "🌹",
+            "Snow": "❄️",
+            "Spacesuit Girl": "👩‍🚀",
+            "Steampunk": "⚙️",
+            "Succubus": "😈",
+            "Sunlight": "☀️",
+            "Weird art": "🎭",
+            "White Hair": "👱‍♀️",
+            "Wings art": "👼",
+            "Woman with Sword": "⚔️"
+        }
 
-    col1, col2, col3 = st.columns(3)
-    
-    for idx, (prompt, emoji) in enumerate(prompt_emojis.items()):
-        full_prompt = f"QT {prompt}"
-        col = [col1, col2, col3][idx % 3]
+        col1, col2, col3 = st.columns(3)
         
-        with col:
-            if st.button(f"{emoji} {prompt}", key=f"btn_{idx}"):
-                handle_prompt_click(full_prompt, idx)
-
-    st.markdown("---")
-    st.markdown("### Generated Artwork:")
-    
-    for key in st.session_state:
-        if key.startswith('selected_prompt_'):
-            idx = key.split('_')[-1]
-            images_key = f'generated_images_{idx}'
+        for idx, (prompt, emoji) in enumerate(prompt_emojis.items()):
+            full_prompt = f"QT {prompt}"
+            col = [col1, col2, col3][idx % 3]
             
-            if images_key in st.session_state:
-                st.write("Prompt:", st.session_state[key])
-                
-                cols = st.columns(len(st.session_state[images_key]))
+            with col:
+                if st.button(f"{emoji} {prompt}", key=f"btn_{idx}"):
+                    handle_prompt_click(full_prompt, idx)
+
+        st.markdown("---")
+        st.markdown("### Generated Artwork:")
+        
+        for key in st.session_state:
+            if key.startswith('selected_prompt_'):
+                idx = key.split('_')[-1]
+                images_key = f'generated_images_{idx}'
                 
-                for col, (model_name, result) in zip(cols, st.session_state[images_key]):
-                    with col:
-                        st.markdown(f"**{model_name}**")
-                        if isinstance(result, str) and result.startswith("Error"):
-                            st.error(result)
-                        else:
-                            # Updated to use use_container_width instead of use_column_width
-                            st.image(result, use_container_width=True)
+                if images_key in st.session_state:
+                    st.write("Prompt:", st.session_state[key])
+                    
+                    cols = st.columns(len(st.session_state[images_key]))
+                    
+                    for col, (model_name, result) in zip(cols, st.session_state[images_key]):
+                        with col:
+                            st.markdown(f"**{model_name}**")
+                            if isinstance(result, str) and result.startswith("Error"):
+                                st.error(result)
+                            else:
+                                st.image(result, use_container_width=True)
+    else:
+        st.info("Please login with your HuggingFace account to use the app")
 
 if __name__ == "__main__":
     main()