import gradio as gr
from PIL import Image
import requests
import io
import base64
import threading

# List of models
models = [
    "dreamlike-art/dreamlike-photoreal-2.0",
    "stabilityai/stable-diffusion-xl-base-1.0",
    "black-forest-labs/FLUX.1-dev",
    "black-forest-labs/FLUX.1-schnell",
    "veryVANYA/ps1-style-flux",
    "alvdansen/softserve_anime",
    "multimodalart/flux-tarot-v1",
    "davisbro/half_illustration",
    "dataautogpt3/OpenDalleV1.1",
    "aleksa-codes/flux-ghibsky-illustration",
    "alvdansen/flux-koda",
    "openskyml/soviet-diffusion-xl",
    "XLabs-AI/flux-RealismLora",
    "alvdansen/frosting_lane_flux",
    "alvdansen/phantasma-anime",
    "kudzueye/Boreal",
    "glif/how2draw",
    "dataautogpt3/FLUX-AestheticAnime",
    "prithivMLmods/Fashion-Hut-Modeling-LoRA",
    "dataautogpt3/FLUX-SyntheticAnime",
    "brushpenbob/flux-midjourney-anime",
    "robert123231/coloringbookgenerator",
    "prithivMLmods/Castor-Collage-Dim-Flux-LoRA",
    "prithivMLmods/Flux-Product-Ad-Backdrop",
    "multimodalart/product-design",
    "glif/90s-anime-art",
    "glif/Brain-Melt-Acid-Art",
    "lustlyai/Flux_Lustly.ai_Uncensored_nsfw_v1",
    "Keltezaa/NSFW_MASTER_FLUX",
    "tryonlabs/FLUX.1-dev-LoRA-Outfit-Generator",
    "Jovie/Midjourney",
    "Yntec/DreamPhotoGASM",
    "strangerzonehf/Flux-Super-Realism-LoRA",
    "stabilityai/stable-diffusion-2-1-base",
    "stabilityai/stable-diffusion-3.5-large",
    "stabilityai/stable-diffusion-3.5-large-turbo",
    "stabilityai/stable-diffusion-3-medium-diffusers",
    "stablediffusionapi/duchaiten-real3d-nsfw-xl",
    "nerijs/pixel-art-xl",
    "KappaNeuro/character-design",
    "alvdansen/sketchedoutmanga",
    "alvdansen/archfey_anime",
    "alvdansen/lofi-cuties",
    "Yntec/YiffyMix",
    "digiplay/AnalogMadness-realistic-model-v7",
    "artificialguybr/selfiephotographyredmond-selfie-photography-lora-for-sdxl",
    "artificialguybr/filmgrain-redmond-filmgrain-lora-for-sdxl",
    "goofyai/Leonardo_Ai_Style_Illustration",
    "goofyai/cyborg_style_xl",
    "alvdansen/littletinies",
    "Dremmar/nsfw-xl",
    "artificialguybr/analogredmond",
    "artificialguybr/PixelArtRedmond",
    "CiroN2022/ascii-art",
    "Yntec/Analog",
    "Yntec/MapleSyrup",
    "digiplay/perfectLewdFantasy_v1.01",
    "digiplay/AbsoluteReality_v1.8.1",
    "goofyai/disney_style_xl",
    "artificialguybr/LogoRedmond-LogoLoraForSDXL-V2",
    "Yntec/epiCPhotoGasm",
]

# Load models
loaded_models = []
available_models = []

for model in models:
    try:
        loaded_models.append(gr.load(f'models/{model}'))
        available_models.append(model)
    except Exception as e:
        print(f"Model {model} could not be loaded: {e}")

print(loaded_models)
print(available_models)

def generate_image(prompt, model_index):
    try:
        model = loaded_models[model_index]
        out_img = model(prompt)
        print(out_img)
        return out_img
    except Exception as e:
        print(e)
        return None

def run_all_models(prompt):
    out_box = []
    out_html = ""

    for i, model in enumerate(loaded_models):
        try:
            out_img = generate_image(prompt, i)
            if out_img:
                raw = Image.open(out_img)
                raw = raw.convert('RGB')
                out_box.append(raw)

                # Convert image to base64 for HTML display
                buffered = io.BytesIO()
                raw.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode()
                img_tag = f"<img src='data:image/png;base64,{img_str}'/>"
                out_html += f"<div class='img_class'><a href='https://huggingface.co/models/{available_models[i]}'>{available_models[i]}</a><br>{img_tag}</div>"
        except Exception as e:
            print(e)
            out_html += f"<div class='img_class'>{available_models[i]}: {str(e)}</div>"

    html_out = f"<div class='grid_class'>{out_html}</div>"
    return out_box, html_out

css = """
.grid_class {
    display: flex;
    flex-wrap: wrap;
    gap: 10px;
}
.img_class {
    min-width: 200px;
    margin: 5px;
}
"""

with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme") as app:
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="Prompt")
            btn = gr.Button("Generate Images")

    out_html = gr.HTML()
    outp = gr.Gallery()
    btn.click(run_all_models, [inp], [outp, out_html])

app.launch()