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

# List of models to process
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",
]

# Initialize a dictionary to track models and their status
model_results = {"success": {}, "failed": []}

def generate_images(prompt):
    output_images = []
    failed_models = []
    for model_name in models:
        try:
            # Attempt to load and generate an image with the model
            model = gr.Interface.load(f"models/{model_name}")
            img_path = model.predict(prompt)

            # Fetch image
            response = requests.get(img_path, stream=True)
            if response.status_code == 200:
                img = Image.open(io.BytesIO(response.content))
                output_images.append((model_name, img))
            else:
                failed_models.append(model_name)
        except Exception as e:
            print(f"Error with model {model_name}: {e}")
            failed_models.append(model_name)

    return output_images, failed_models

def app_interface(prompt):
    output_images, failed_models = generate_images(prompt)
    images_html = ""
    for model_name, img in output_images:
        # Convert images for display in Gradio
        img_buffer = io.BytesIO()
        img.save(img_buffer, format="PNG")
        img_str = base64.b64encode(img_buffer.getvalue()).decode()
        images_html += f"""
        <div>
            <strong>{model_name}</strong>
            <img src='data:image/png;base64,{img_str}' style="max-width: 200px;" />
        </div>
        """
    
    failed_html = "<ul>" + "".join(f"<li>{model}</li>" for model in failed_models) + "</ul>"
    return images_html, f"<h3>Failed Models:</h3>{failed_html}"

# Define Gradio UI
with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt")
            generate_button = gr.Button("Generate Images")
        with gr.Column():
            output_html = gr.HTML(label="Generated Images")
            failed_html = gr.HTML(label="Failed Models")

    generate_button.click(app_interface, inputs=[prompt_input], outputs=[output_html, failed_html])

app.launch()