import gradio as gr
from gradio_client import Client, handle_file
from gradio_imageslider import ImageSlider

tile_upscaler_url = "gokaygokay/TileUpscalerV2"

client_tile_upscaler = Client(tile_upscaler_url)

def gradio_process_image(image, resolution, num_inference_steps, strength, hdr, guidance_scale, controlnet_strength, scheduler_name):
  result = client_tile_upscaler.predict(
      param_0=handle_file(image),
      param_1=resolution,
      param_2=num_inference_steps,
      param_3=strength,
      param_4=hdr,
      param_5=guidance_scale,
      param_6=controlnet_strength,
      param_7=scheduler_name,
      api_name="/wrapper"
  )
  return result


with gr.Blocks() as demo:
  with gr.Row():
    with gr.Column():
      input_image = gr.Image(type="pil", label="Input Image")
      run_button = gr.Button("Enhance Image")
    with gr.Column():
      output_image = gr.Image(type="pil", label="Output Image")
      # output_slider = ImageSlider(label="Before / After", type="numpy")

  with gr.Accordion("Advanced Options", open=False):
    resolution = gr.Slider(minimum=128, maximum=2048, value=1024, step=128, label="Resolution")
    num_inference_steps = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Number of Inference Steps")
    strength = gr.Slider(minimum=0, maximum=1, value=0.2, step=0.01, label="Strength")
    hdr = gr.Slider(minimum=0, maximum=1, value=0, step=0.1, label="HDR Effect")
    guidance_scale = gr.Slider(minimum=0, maximum=20, value=6, step=0.5, label="Guidance Scale")
    controlnet_strength = gr.Slider(minimum=0.0, maximum=2.0, value=0.75, step=0.05, label="ControlNet Strength")
    scheduler_name = gr.Dropdown(
        choices=["DDIM", "DPM++ 3M SDE Karras", "DPM++ 3M Karras"],
        value="DDIM",
        label="Scheduler"
    )


    run_button.click(
      fn=gradio_process_image, 
      inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale, controlnet_strength, scheduler_name],
      outputs=output_image
    )
    
demo.launch()