import gradio as gr
from ModelLoader import ModelLoader
import torch
import os

max_img_wh = 4096  # Set the maximum image size in pixels

def inference(image, use_gpu):
    gpu_ids = '0' if use_gpu else ''
    model = ModelLoader(gpu_ids=gpu_ids, max_img_wh=max_img_wh)
    model.load()
    output_img = model.inference(image_pil=image)
    return output_img

demo = gr.Interface(
    fn=inference,
    inputs=[
        gr.components.Image(type='pil', label='Input Image'),
        gr.components.Checkbox(label='Use GPU', value=torch.cuda.is_available())  # Precheck the GPU checkbox
    ],
    outputs=gr.components.Image(type='pil', label='Output Image', format='jpeg'),
    title="Underwater Photo Enhancement",
    description="This tool enhances the quality of underwater photos by improving clarity and color. Upload your underwater photo to see the enhanced version.",
    flagging_mode='never'
)
demo.launch()