import gradio as gr
from gradio_client import Client
import numpy as np

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

flux_1_schell_spaces = [
    "https://black-forest-labs-flux-1-schnell.hf.space",
    "ChristianHappy/FLUX.1-schnell",
    "innoai/FLUX.1-schnell",
    "tuan2308/FLUX.1-schnell",
    "FiditeNemini/FLUX.1-schnell"
]

client = None
job = None

def infer(selected_space, prompt, seed=42, randomize_seed=False, width=1024, height=1024, 
          num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    global job
    global client
    
    if client is None:
        try:
            client = Client(selected_space)
            print(f"Loaded custom model from {selected_space}")
        except ValueError as e:
            client = None
            print(f"Failed to load custom model from {selected_space}: {e}")
            raise gr.Error("Failed to load client after trying all spaces.")
            
    try:
        job = client.submit(
            prompt=prompt,
            seed=seed,
            randomize_seed=randomize_seed,
            width=width,
            height=height,
            num_inference_steps=num_inference_steps,
            api_name="/infer"
        )
        result = job.result()
    except ValueError as e:
        client = None
        raise gr.Error(e)
    return result


css = """
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    selected_space_index = gr.State(0)
    
    with gr.Column(elem_id="col-container"):
        space = gr.Radio(
            flux_1_schell_spaces,
            label="Choose Your Model",
            value=flux_1_schell_spaces[0]
        )
        
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
                value="beautiful sexy woman, at the beach, relaxing, dreamy, highly detailed, big boobs, realistic lighting, sharp focus"
            )
            run_button = gr.Button("Run", scale=0)
            
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
        
        gr.on(
            triggers=[run_button.click, prompt.submit],
            fn=infer,
            inputs=[
                space,
                prompt,
                seed,
                randomize_seed,
                width,
                height,
                num_inference_steps
            ],
            outputs=[result, seed]
        )

        with gr.Row():
            gr.Markdown("""
            **Best AI Tools**
            [Nude AI](https://nudeai.beauty) • [Uncensored AI](https://uncensoredai.cc) • [AI Hentai Generator](https://aihentaigenerator.fun) • [Stable Diffusion Hentai](https://stable-diffusion-hentai.aihentaigenerator.fun) • [Bing Image Creator](https://bingimagecreator.online) • [NSFW AI Art](https://nsfwaiart.art) • [NSFW AI Chatbot](https://nsfw-ai-chatbot.online) • [NSFW AI World](https://nsfwai.world) • [AI Tools Directory](https://aitoolsdirectory.online) • [Viggle AI](https://viggleai.live) • [TDEE Calculator](https://tdeecalculator.online) • [Calculator App](https://calculatorapp.online) • [Compound Interest Calculator](https://compoundinterestcalculator.site) • [AI Story Generator](https://aistorygenerator.fun) • [Llama AI](https://llamaai.online) • [AI Art Free](https://aiartfree.online) • [Cek Khodam](https://cekkhodam.co) • [AI Detector](https://ai-detector.online) • [AI Cover](https://aicover.fun) • [Flux AI](https://flux-ai.online) • [Brat Generator](https://bratgenerator.org) • [AI Text Generator](https://aitextgenerator.live) • [Black Myth Wukong](https://blackmythwukong.buzz) • [Book Summary](https://booksummary.wiki) • [ChatGPT Gratuit](https://chatgptgratuit.chat) • [ChatGPT Online](https://chatgtponline.com) • [Gau Thai](https://gauthai.pro) • [Best Recipe](https://bestrecipe.pro)
            **Games**
            [Cookie Clicker Unblocked](https://cookie-clicker-unblocked.pro) • [Freaky Font](https://freakyfont.org) • [Chat Gratuito](https://chatgratuito.online) • [Incredibox Sprunki](https://incrediboxsprunki.pro) • [Moto X3M Unblocked](https://motox3munblocked.pro) • [Parkour Civilization](https://parkourcivilization.pro) • [Sled Rider 3D](https://sledrider3d.com) • [Snow Rider 3D](https://snow-rider-3d.pro) • [Sprunki Game](https://sprunkigame.pro) • [Stickman Hook Unblocked](https://stickmanhookunblocked.pro) • [Escape Road](https://escaperoad.fun) • [Sprunked](https://sprunked.pro) • [Abgerny](https://abgerny.buzz)• [Sprunki Phase](https://sprunki-phase.online) • [Block Blast](https://block-blast.pro) • [Traffic Racer](https://trafficracer.online) • [Football Bros](https://footballbros.online) • [Sprunki Parasite](https://sprunkiparasite.com) [Sprunki Mustard](https://sprunki-mustard.online) • [PVZ Fusion Download](https://pvz-fusion.online) • [My Femboy Roommate](https://myfemboy-roommate.com) • [Homicipher](https://homicipher.fun) • [Sprunki Parasite](https://sprunkiparasite.com) • [Sprunki Retake](https://sprunkiretake.online) • [Sprunki Mustard](https://sprunki-mustard.online) • [PVZ Fusion Download](https://pvz-fusion.online) • [My Femboy Roommate](https://myfemboy-roommate.com) • [Homicipher](https://homicipher.fun) • [Anime Reborn Codes](https://animereborncodes.com) • [The Rise of the Golden Idol](https://theriseofthegoldenidol.com/) • [Five Hearts Under One Roof](https://fiveheartsunderoneroof.com) • [Googlies](https://googlies.pro/)
            """)

demo.launch(share=True)