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

tile_upscaler_url = "https://gokaygokay-tileupscalerv2.hf.space"

client_tile_upscaler = Client(tile_upscaler_url)

def upscale_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

def clear_output(image_slider):
  image_slider[0] = None
  image_slider[1] = None
  return image_slider
  

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

  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=clear_output,
      inputs=[output_slider],
      outputs=[output_slider]
    ).then(
      upscale_image, 
      [input_image, resolution, num_inference_steps, strength, hdr, guidance_scale, controlnet_strength, scheduler_name],
      output_slider
    )
    
demo.launch()