import torch
from PIL import Image
import gradio as gr
import spaces
from KandiSuperRes import get_SR_pipeline

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
sr_pipe2x = get_SR_pipeline(device=device, fp16=True, flash=True, scale=2)
sr_pipe4x = get_SR_pipeline(device=device, fp16=True, flash=False, scale=4)


@spaces.GPU()
def inference(image, size):
    if image is None:
        raise gr.Error("Image not uploaded")
    # r_image = Image.open(image)
    
    if size == '2x':
        result = sr_pipe2x(image)
    else:
        result = sr_pipe4x(image)

    print(f"Image size ({device}): {size} ... OK")
    
    return result


title = "KandiSuperRes - diffusion model for super resolution"
description = "KandiSuperRes Flash is a new version of the diffusion model for super resolution. This model includes a distilled version of the KandiSuperRes model and a distilled model Kandinsky 3.0 Flash. KandiSuperRes Flash not only improves image clarity, but also corrects artifacts, draws details, improves image aesthetics. And one of the most important advantages is the ability to use the model in the «infinite super resolution» mode."
article = "<div style='text-align: center;'>Twitter <a href='https://twitter.com/DoEvent' target='_blank'>Max Skobeev</a> | <a href='https://huggingface.co/ai-forever/KandiSuperRes' target='_blank'>Model card</a><div>"


gr.Interface(inference,
    [gr.Image(type="pil"), 
    gr.Radio(['2x', '4x'], 
    type="value",
    value='2x',
    label='Resolution model')], 
    gr.Image(type="filepath", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[['groot.jpeg', "2x"]],
    allow_flagging='never',
    cache_examples=False,
    delete_cache=(1800, 3600),
    ).queue(api_open=True).launch(show_error=True, show_api=True)