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

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

class ModelGenerator:
    @staticmethod
    def generate_midjourney(prompt):
        try:
            client = Client("mukaist/Midjourney", hf_token=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",
                use_negative_prompt=True,
                style="2560 x 1440",
                seed=0,
                width=1024,
                height=1024,
                guidance_scale=6,
                randomize_seed=True,
                api_name="/run"
            )
            
            # Handle different types of results
            if isinstance(result, tuple):
                # If it's a tuple, the first element might be the image gallery
                if len(result) > 0 and isinstance(result[0], list):
                    image_data = result[0][0]  # Get first image from gallery
                    if isinstance(image_data, dict) and 'image' in image_data:
                        return ("Midjourney", image_data['image'])
                    elif isinstance(image_data, (str, bytes)):
                        return ("Midjourney", image_data)
                else:
                    return ("Midjourney", result[0])  # Try first element of tuple
            elif isinstance(result, list) and len(result) > 0:
                # If it's a list, get the first element
                image_data = result[0]
                if isinstance(image_data, dict) and 'image' in image_data:
                    return ("Midjourney", image_data['image'])
                else:
                    return ("Midjourney", image_data)
            elif isinstance(result, str):
                # If it's a direct string (URL or path)
                return ("Midjourney", result)
            else:
                return ("Midjourney", f"Error: Unexpected result format: {type(result)}")
        except Exception as e:
            return ("Midjourney", f"Error: {str(e)}")

    @staticmethod
    def generate_stable_cascade(prompt):
        try:
            client = Client("multimodalart/stable-cascade", hf_token=HF_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"
            )
            return ("Stable Cascade", result)
        except Exception as e:
            return ("Stable Cascade", f"Error: {str(e)}")

    @staticmethod
    def generate_stable_diffusion_3(prompt):
        try:
            client = Client("stabilityai/stable-diffusion-3-medium", hf_token=HF_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"
            )
            return ("SD 3 Medium", result)
        except Exception as e:
            return ("SD 3 Medium", f"Error: {str(e)}")

    @staticmethod
    def generate_stable_diffusion_35(prompt):
        try:
            client = Client("stabilityai/stable-diffusion-3.5-large", hf_token=HF_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"
            )
            return ("SD 3.5 Large", result)
        except Exception as e:
            return ("SD 3.5 Large", f"Error: {str(e)}")

    @staticmethod
    def generate_playground_v2_5(prompt):
        try:
            client = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/ji5gy/", hf_token=HF_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"
            )
            # 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):
    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
        }
        
        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):
            results.append(future.result())
    
    return results

def handle_prompt_click(prompt_text, key):
    if not HF_TOKEN:
        st.error("Environment variable 'ArtToken' is not set!")
        return
        
    st.session_state[f'selected_prompt_{key}'] = prompt_text
    
    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[f'generated_images_{key}'] = results
        st.success("Artwork generated successfully!")

def main():
    st.title("🎨 Multi-Model Art Generator")

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

    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:")
    
    for key in st.session_state:
        if key.startswith('selected_prompt_'):
            idx = key.split('_')[-1]
            images_key = f'generated_images_{idx}'
            
            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:
                            # Updated to use use_container_width instead of use_column_width
                            st.image(result, use_container_width=True)

if __name__ == "__main__":
    main()