import gradio as gr
from all_models import models
from externalmod import gr_Interface_load
import asyncio
import os
from datetime import datetime

# Load the models
HF_TOKEN = os.getenv("HF_TOKEN", None)
from PIL import Image
import io


def load_models(models):
    loaded_models = {}
    for model in models:
        try:
            loaded_models[model] = gr_Interface_load(f'models/{model}', hf_token=HF_TOKEN)
        except Exception as e:
            print(f"Error loading {model}: {e}")
    return loaded_models

models_load = load_models(models)

# Generate image function
async def infer(model_str, prompt, seed=-1):
    task = asyncio.create_task(
        asyncio.to_thread(models_load[model_str].fn, prompt=prompt, seed=seed, token=HF_TOKEN)
    )
    await asyncio.sleep(0)
    result = await asyncio.wait_for(task, timeout=600)
    return result


def generate_image(model_name, prompt, seed):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        # Get the result from inference
        result = loop.run_until_complete(infer(model_name, prompt, seed))
        if isinstance(result, tuple):
            # Assuming the first element is the image data
            result = result[0]
        if isinstance(result, bytes):
            # Convert bytes to PIL Image if necessary
            return Image.open(io.BytesIO(result))
        elif isinstance(result, Image.Image):
            return result
        else:
            raise ValueError(f"Unexpected output type: {type(result)}")
    finally:
        loop.close()

        
# Interface
with gr.Blocks() as demo:
    with gr.Column():
        model_choice = gr.Dropdown(
            choices=models, label="Select Model", value=models[0]
        )
        prompt_input = gr.Textbox(label="Enter your prompt")
        seed_input = gr.Slider(
            label="Seed", minimum=-1, maximum=100000, step=1, value=-1
        )
        generate_button = gr.Button("Generate Image")
        output_image = gr.Image(label="Generated Image", show_download_button=True)

        generate_button.click(
            fn=generate_image,
            inputs=[model_choice, prompt_input, seed_input],
            outputs=[output_image],
        )

demo.launch()