import gradio as gr
import numpy as np
import random
import time
from optimum.intel import OVStableDiffusionXLPipeline
import torch
from diffusers import EulerDiscreteScheduler
from io import BytesIO
from PIL import Image
import base64

model_id = "None1145/noobai-XL-Vpred-0.65s-openvino"

prev_height = 1216
prev_width = 832

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

def reload_model(new_model_id):
    global pipe, model_id, prev_height, prev_width
    model_id = new_model_id
    try:
        print(f"{model_id}...")
        pipe = OVStableDiffusionXLPipeline.from_pretrained(model_id, compile=False)
        if model_id == "None1145/noobai-XL-Vpred-0.65s-openvino":
            scheduler_args = {"prediction_type": "v_prediction", "rescale_betas_zero_snr": True}
            pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, **scheduler_args)
        pipe.reshape(batch_size=1, height=prev_height, width=prev_width, num_images_per_prompt=1)
        pipe.compile()
        print(f"{model_id}!!!")
        return f"Model successfully loaded: {model_id}"
    except Exception as e:
        return f"Failed to load model: {str(e)}"
reload_model(model_id)

def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
):
    global prev_width, prev_height, pipe

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    if prev_width != width or prev_height != height:
        pipe.reshape(batch_size=1, height=height, width=width, num_images_per_prompt=1)
        pipe.compile()
        prev_width = width
        prev_height = height

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image, seed


examples = ["murasame \(senren\), senren banka",]
with gr.Blocks() as img:
    gr.Markdown("# OpenVINO Text to Image")
    gr.Markdown("### It usually takes 2200 seconds to generate an 832x1216 image (28 steps) (CPU).")
    
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
                value="murasame \(senren\), senren banka"
            )

            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=1,
                maximum=60,
                step=1,
                value=28,
            )

            run_button = gr.Button("Run", scale=0, variant="primary")

        result = gr.Image(label="Result", show_label=False, value=Image.open(requests.get("https://huggingface.co/None1145/noobai-XL-Vpred-0.65s-openvino/blob/main/example.webp").content))

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=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=512,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=832,
                )

                height = gr.Slider(
                    label="Height",
                    minimum=512,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1216,
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=5.0,
                )

        gr.Examples(examples=examples, inputs=[prompt])

    gr.Markdown("### Model Reload")
    with gr.Row():
        new_model_id = gr.Text(label="New Model ID", placeholder="Enter model ID", value=model_id)
        reload_button = gr.Button("Reload Model", variant="primary")
        reload_status = gr.Text(label="Status", interactive=False)

    reload_button.click(
        fn=reload_model,
        inputs=new_model_id,
        outputs=reload_status,
    )

    run_button.click(
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    img.queue(max_size=10).launch()