import gradio as gr
from stablepy import Model_Diffusers


# Define the function to generate images
def generate_image(prompt, num_steps, guidance_scale, sampler, img_width, img_height, upscaler_model_path, upscaler_increases_size, hires_steps, base_model_id):

    model = Model_Diffusers(
        base_model_id=base_model_id,
        task_name='txt2img',
    )

    image, info_image = model(
        prompt=prompt,
        num_steps=num_steps,
        guidance_scale=guidance_scale,
        sampler=sampler,
        img_width=img_width,
        img_height=img_height,
        upscaler_model_path=upscaler_model_path,
        upscaler_increases_size=upscaler_increases_size,
        hires_steps=hires_steps,
    )
    return image[0]

# Create the Gradio Blocks UI
with gr.Blocks(gr.themes.Ocean(), title="StablePY") as demo:
    gr.Markdown("# StablePY")

    moelpah = gr.Textbox(label="model", placeholder="user/repo")
    
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", value="highly detailed portrait of an underwater city, with towering spires and domes rising up from the ocean floor")
    
    with gr.Row():
        num_steps = gr.Slider(label="Number of Steps", minimum=1, maximum=100, value=30)
        guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=7.5)
        sampler = gr.Dropdown(label="Sampler", choices=["DPM++ 2M", "DDIM", "Euler A"], value="DPM++ 2M")
    
    with gr.Row():
        img_width = gr.Number(label="Image Width", value=1024)
        img_height = gr.Number(label="Image Height", value=1024)
    
    with gr.Row():
        upscaler_model_path = gr.Textbox(label="Upscaler Model Path", value="./upscaler/RealESRGAN_x4plus_anime_6B.pth", visible=False)
        upscaler_increases_size = gr.Number(label="Upscaler Increase Size", value=1.5)
        hires_steps = gr.Number(label="Hires Steps", value=25)

    with gr.Row():
        generate_button = gr.Button("Generate")
    
    with gr.Row():
        output_image = gr.Image(label="Generated Image")

    # Link the button to the generation function
    generate_button.click(
        generate_image,
        inputs=[moelpah, prompt, num_steps, guidance_scale, sampler, img_width, img_height, upscaler_model_path, upscaler_increases_size, hires_steps],
        outputs=[output_image]
    )

# Launch the app
demo.launch()