import os
import warnings
from pathlib import Path

import gradio as gr
from gradio_imageslider import ImageSlider
import numpy as np
import torch

import fastai
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
from huggingface_hub import snapshot_download

os.system("pip freeze")
from collections.abc import Sized # Import Sized from collections.abc
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")


# private repo
token = os.getenv("HF_TOKEN")
repo_id = "afondiel/image-colorizer-deoldify"
repo_type = "space"

# Download the snapshot from the space repository
snapshot_folder = snapshot_download(repo_id=repo_id, repo_type=repo_type, token=token)

# Set the device to use for computation
# choices:  CPU, GPU0...GPU7
if torch.cuda.is_available():
    device.set(device=DeviceId.GPU0)
else:
    device.set(device=DeviceId.CPU)

# Load the pre-trained model
_colorizer = get_image_colorizer(root_folder=Path(snapshot_folder), artistic=True)


def colorizer_fn(input_img, render_factor):
  """
  Colorize grayscale images/photos
  - @param input_img old (grayscale) image 
  - @param render_factor render_factor
  """
  if input_img is not None and input_img !='':
    output_img = _colorizer.get_transformed_image(
        path=input_img,
        render_factor=int(render_factor),
        watermarked=watermarked,
        post_process=True,
    )
  else:
    print('Provide an image and try again.')

  return (input_img, output_img) # Return a tuple of old and color Image to be plotted with ImageSlider()

title = "AI Image Colorizer"
description = "Colorize old images with AI"
examples = [["./demo.jpg"],]

demo = gr.Interface(
    fn=colorizer_fn,
    inputs=[gr.Image(type="filepath" , label="Old image"), gr.Slider(0, 40, label="Render Factor", value=10)],
    outputs=ImageSlider(type="pil", label="Old vs Colored image"),
    examples=examples,
    title=title,
    description=description,
)

# Launch the demo

if __name__ == "__main__":
  demo.launch()