import gradio as gr
import torch
from bsrgan import BSRGAN

# Images
torch.hub.download_url_to_file('https://raw.githubusercontent.com/kadirnar/bsrgan-pip/main/data/images/butterfly.png', 'butterfly.jpg')

def bsrgan_inference(
    image: gr.inputs.Image = None,
    model_path: gr.inputs.Dropdown = 'kadirnar/bsrgan',
):
    """
    BSRGAN inference function
    Args:
        image: Input image
        model_path: Path to the model
    Returns:
        Rendered image
    """
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    model = BSRGAN(model_path, device=device, hf_model=True)
    pred = model.predict(img_path=image)
    return pred
        

inputs = [
    gr.inputs.Image(type="filepath", label="Input Image"),
    gr.inputs.Dropdown(
        label="Model",
        choices=[
        "kadirnar/bsrgan",
        "kadirnar/BSRGANx2",
        "kadirnar/RRDB_PSNR_x4",
        "kadirnar/RRDB_ESRGAN_x4",
        "kadirnar/DF2K",
        "kadirnar/DPED",
        "kadirnar/DF2K_JPEG",
        ],
        default="kadirnar/bsrgan",
    ),
]

outputs = gr.outputs.Image(type="filepath", label="Output Image")
title = "BSRGAN: Designing a Practical Degradation Model for Deep Blind Image Super-Resolution."
description = "BSRGAN for Deep Blind Image Super-Resolution model aims to design a practical degradation model for deep blind image super-resolution by considering the deterioration of image quality over time. It uses deep learning methods to predict the deterioration of image quality and to assist in the re-creation of images at higher resolution using these predictions."
examples = [["butterfly.jpg", "kadirnar/bsrgan"]]

demo_app = gr.Interface(
    fn=bsrgan_inference,
    inputs=inputs,
    outputs=outputs,
    title=title,
    description=description,
    examples=examples,
    cache_examples=True,
    live=True,
    theme='huggingface',
)
demo_app.launch(debug=True, enable_queue=True)