import gradio as gr
import time
from openvino.runtime import Core
from optimum.intel import OVStableDiffusionPipeline
from PIL import Image as PImg
import numpy as np

model_id = "NoCrypt/SomethingV2_2"
core = Core()

# Check available devices
devices = core.available_devices
print("Available devices:", devices)

# Use 'GPU' if available, otherwise fall back to 'CPU'
device = "GPU" if "GPU" in devices else "CPU"
print(f"Using device: {device}")

ov_pipe_bf16 = OVStableDiffusionPipeline.from_pretrained(
    model_id,
    export=True,
    device=device  # Specify the device here
)

# The compile step is not needed as it's handled internally
def generate_image(prompt, num_inference_steps):
    # Generate an image from the prompt using the pipeline
    start = time.time()
    output = ov_pipe_bf16(prompt, num_inference_steps=num_inference_steps, output_type="np")
    end = time.time()
    print("Inference time: ", end - start)
    
    # Convert the image to PIL Image object
    image_data = output.images[0]
    image_data = (image_data * 255).clip(0, 255).astype(np.uint8)
    image = PImg.fromarray(image_data)
    
    # Calculate the target size based on the scaling factor
    target_resolution = 1.2
    width = int(image.width * target_resolution)
    height = int(image.height * target_resolution)
    target_size = (width, height)
    
    # Upscale the image to the target resolution
    upscaled_image = image.resize(target_size, resample=PImg.BICUBIC)
    return upscaled_image


examples = [
    ["masterpiece, best quality, 1girl, blonde, colorful, clouds, outdoors, falling leaves, smiling, whimsical"],
    ["masterpiece, best quality, landscape"],
    ["masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt"]
]

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter a prompt"),
        gr.Slider(minimum=1, maximum=20, value=8, step=1, label="Number of inference steps")
    ],
    outputs=gr.Image(label="Generated image"),
    title="OpenVINO Anime Diffusion",
    description="A gradio app that generates an image from a text prompt using the stable diffusion pipeline using the OpenVINO library for speed!",
    examples=examples,
)

iface.launch()