import gradio as gr
import torch
from optimum.intel.openvino.modeling_diffusion import OVModelVaeDecoder, OVBaseModel, OVStableDiffusionPipeline
from huggingface_hub import snapshot_download
import openvino.runtime as ov
from typing import Optional, Dict

model_id = "hsuwill000/Fluently-v4-LCM-openvino"

HIGH = 1024
WIDTH = 512

batch_size = -1  # Or set it to a specific positive integer if needed

class CustomOVModelVaeDecoder(OVModelVaeDecoder):
    def __init__(
        self, model: ov.Model, parent_model: OVBaseModel, ov_config: Optional[Dict[str, str]] = None, model_dir: str = None,
    ):
        super(OVModelVaeDecoder, self).__init__(model, parent_model, ov_config, "vae_decoder", model_dir)



pipe = OVStableDiffusionPipeline.from_pretrained(
    model_id,
    compile=False,
    ov_config={"CACHE_DIR": ""},
    torch_dtype=torch.bfloat16,  # More standard dtype for speed
    safety_checker=None,
    use_safetensors=False,
)

taesd_dir = snapshot_download(repo_id="deinferno/taesd-openvino")

pipe.vae_decoder = CustomOVModelVaeDecoder(model = OVBaseModel.load_model(f"{taesd_dir}/vae_decoder/openvino_model.xml"), 
                                           parent_model = pipe, 
                                           model_dir = taesd_dir
                                          )

print(pipe.scheduler.compatibles)

pipe.reshape(batch_size=batch_size, height=HIGH, width=WIDTH, num_images_per_prompt=1)

pipe.compile()

prompt = ""
negative_prompt = "Easy Negative, worst quality, low quality, normal quality, lowers, monochrome, grayscales, skin spots, acnes, skin blemishes, age spot, 6 more fingers on one hand, deformity, bad legs, error legs, bad feet, malformed limbs, extra limbs, ugly, poorly drawn hands, poorly drawn feet, poorly drawn face, text, mutilated, extra fingers, mutated hands, mutation, bad anatomy, cloned face, disfigured, fused fingers"

def infer(prompt, negative_prompt, num_inference_steps=8):
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=WIDTH,
        height=HIGH,
        guidance_scale=1.0,
        num_inference_steps=num_inference_steps,
        num_images_per_prompt=1,
    ).images[0]
    
    return image

css = """
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

power_device = "CPU"

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # {model_id.split('/')[1]} {WIDTH}x{HIGH}
        Currently running on {power_device}.
        """)
        
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )         
            run_button = gr.Button("Run", scale=1)
        
        result = gr.Image(label="Result", show_label=False)

    run_button.click(
        fn=infer,
        inputs=[prompt],
        outputs=[result]
    )

demo.queue().launch()