backup1-app.py

import streamlit as st
from gradio_client import Client
import time
import concurrent.futures
import os
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")

class ModelGenerator:
    def __init__(self, token):
        self.token = token

    def generate_midjourney(self, prompt):
        try:
            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",
                use_negative_prompt=True,
                style="2560 x 1440",
                seed=0,
                width=1024,
                height=1024,
                guidance_scale=6,
                randomize_seed=True,
                api_name="/run"
            )
            
            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 ("Midjourney", Image.open(io.BytesIO(response.content)))
                    else:
                        return ("Midjourney", Image.open(image_data))
                else:
                    return ("Midjourney", Image.open(io.BytesIO(image_data)))
            return ("Midjourney", "Error: No image generated")
        except Exception as e:
            return ("Midjourney", f"Error: {str(e)}")

    def generate_stable_cascade(self, prompt):
        try:
            client = Client("multimodalart/stable-cascade", hf_token=self.token)
            result = client.predict(
                prompt=prompt,
                negative_prompt=prompt,
                seed=0,
                width=1024,
                height=1024,
                prior_num_inference_steps=20,
                prior_guidance_scale=4,
                decoder_num_inference_steps=10,
                decoder_guidance_scale=0,
                num_images_per_prompt=1,
                api_name="/run"
            )
            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)}")

    def generate_stable_diffusion_3(self, prompt):
        try:
            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=self.token)
            result = client.predict(
                prompt=prompt,
                negative_prompt=prompt,
                seed=0,
                randomize_seed=True,
                width=1024,
                height=1024,
                guidance_scale=5,
                num_inference_steps=28,
                api_name="/infer"
            )
            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)}")

    def generate_stable_diffusion_35(self, prompt):
        try:
            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=self.token)
            result = client.predict(
                prompt=prompt,
                negative_prompt=prompt,
                seed=0,
                randomize_seed=True,
                width=1024,
                height=1024,
                guidance_scale=4.5,
                num_inference_steps=40,
                api_name="/infer"
            )
            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)}")

    def generate_playground_v2_5(self, prompt):
        try:
            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", hf_token=self.token)
            result = client.predict(
                prompt,
                prompt,  # negative prompt
                True,    # use negative prompt
                0,      # seed
                1024,   # width
                1024,   # height
                7.5,    # guidance scale
                True,   # randomize seed
                api_name="/run"
            )
            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):
    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
        }
        
        for model in selected_models:
            if model in model_map:
                futures.append(executor.submit(model_map[model], prompt))
        
        for future in concurrent.futures.as_completed(futures):
            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.is_authenticated:
        st.error("Please login with your HuggingFace account first!")
        return
        
    selected_models = st.session_state.get('selected_models', [])
    
    if not selected_models:
        st.warning("Please select at least one model from the sidebar!")
        return

    with st.spinner('Generating artwork...'):
        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()

    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:")

        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]
            
            with col:
                if st.button(f"{emoji} {prompt}", key=f"btn_{idx}"):
                    handle_prompt_click(full_prompt, idx)

        st.markdown("---")
        st.markdown("### Generated Artwork:")
        
        # Display any previously generated images
        for key in st.session_state:
            if key.startswith('generated_images_'):
                idx = key.split('_')[-1]
                prompt_key = f'selected_prompt_{idx}'
                
                if prompt_key in st.session_state:
                    st.write("Prompt:", st.session_state[prompt_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)
                            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")

if __name__ == "__main__":
    main()