#NOTE:  This must be the first call in order to work properly!
from deoldify import device
from deoldify.device_id import DeviceId
#choices:  CPU, GPU0...GPU7
device.set(device=DeviceId.GPU0)

import torch

if not torch.cuda.is_available():
    print('GPU not available.')
#!mkdir 'models'
#!wget https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth -O ./models/ColorizeArtistic_gen.pth
import fastai
from deoldify.visualize import *
import warnings
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")
colorizer = get_image_colorizer(artistic=True)
source_url = 'https://i.pinimg.com/originals/09/c2/22/09c222cc275efa908bdc140a750ca4e4.jpg'
render_factor = 10
watermarked = False
image_path_1 = colorizer.plot_transformed_image_from_url(url=source_url, render_factor=render_factor, compare=True, watermarked=watermarked)
#image_path_2 = colorizer.plot_transformed_image(path=source_path, render_factor=render_factor, compare=True)
print('image_path',image_path_1)

import gradio as gr
import tensorflow as tf
import numpy as np
import os

def Colorize(address,watermark,URL,render_factor):
    text_box_output = "Successful"
    watermarked = True
    output_address = np.array([
    [0,0],
    [0,0],
], dtype=np.uint8)
    if watermark == "Watermark On":
      watermarked = True
      render_factor = int(render_factor)
    elif watermark == "Watermark Off":
      watermarked = False
      render_factor = int(render_factor)
    else:
      text_box_output = "Please set the watermark setting."
    if URL != '':
      try:
         output_address = colorizer.plot_transformed_image_from_url(url=URL, render_factor=render_factor, compare=True, watermarked=watermarked)
      except:
        text_box_output = 'The URL is invalid'
    else:
      if address == None:
        text_box_output = "Please provide the image"
      else:
       output_address = colorizer.plot_transformed_image(path=address, render_factor=render_factor, compare=True, watermarked=watermarked)
      if address.endswith('tiff') == True:
         text_box_output = "Files with tiff format cannot be displayed properly in this website, please download the image to view"
      else:
         text_box_output = "Colorization Accomplished"
    return text_box_output,output_address
with gr.Blocks() as demo:
    image_box=gr.Image(type="filepath")
    image_output_box = gr.Image()
   # text_box_1 = gr.Textbox(value="10",label="The render factor")
    text_box_2 = gr.Textbox(label='URL')
    drop_down_1 = gr.Dropdown(['Watermark On','Watermark Off'],label='WaterMark')
    slider = gr.Slider(7,40,step=1,label='render factor')
    text_box_output = gr.Textbox(label="Status",interactive=False)
    slider.release(Colorize, inputs=[image_box,drop_down_1,text_box_2,slider], outputs=[text_box_output,image_output_box])
    #b1=gr.Button("Colorized!")
    #b1.click(Colorize, inputs=[image_box,drop_down_1,text_box_2,text_box_1], outputs= image_output_box)

demo.queue().launch(share=True,debug=True)