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 HfApi, snapshot_download
from huggingface_hub import HfApi, HfFolder

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

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

# running primarly on CPU
device.set(device=DeviceId.CPU)

# Initialize Hugging Face API
api = HfApi(token=token)

# Authenticate
HfFolder.save_token(token)

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

# if GPU available
device.set(device=DeviceId.GPU0)

# 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=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(minimum=0, maximum=40, step=1, 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()