import os
import sys
import subprocess
import requests
import gradio as gr
import PIL.Image
import imageio
import numpy
import skimage.transform
from shutil import copyfileobj
from skimage import img_as_ubyte

#Clonar repo si no existe
if not os.path.isdir('first-order-model'):
    subprocess.run(['git', 'clone', 'https://github.com/AliaksandrSiarohin/first-order-model'])

sys.path.insert(0, 'first-order-model')
from demo import load_checkpoints, make_animation

#Descargar checkpoint si no existe (En este caso el de caras)
filename = 'vox-cpk.pth.tar'
if not os.path.isfile(filename):
    response = requests.get(
        'https://github.com/graphemecluster/first-order-model-demo/releases/download/checkpoints/' + filename,
        stream=True
    )
    with open(filename, 'wb') as f:
        copyfileobj(response.raw, f)
    print('Checkpoint descargado')

#Detectar CPU o GPU (en nuestro caso será cpu, pero nunca se sabe)
import torch
use_cpu = not torch.cuda.is_available()
print(f'Usando: {"CPU" if use_cpu else "GPU"}')

generator, kp_detector = load_checkpoints(
    config_path='first-order-model/config/vox-256.yaml',
    checkpoint_path=filename,
    cpu=use_cpu
)
print('Modelo cargado')

def resize(image, size=(256, 256)):
    w, h = image.size
    d = min(w, h)
    r = ((w-d)//2, (h-d)//2, (w+d)//2, (h+d)//2)
    return image.resize(size, resample=PIL.Image.LANCZOS, box=r)

def generate(source_image, driving_video):
    img = PIL.Image.fromarray(source_image).convert("RGB")
    img = resize(img)

    reader = imageio.get_reader(driving_video, mode='I', format='FFMPEG')
    fps = reader.get_meta_data()['fps']
    driving_frames = [frame for frame in reader]

    predictions = make_animation(
        skimage.transform.resize(numpy.asarray(img), (256, 256)),
        [skimage.transform.resize(f, (256, 256)) for f in driving_frames],
        generator,
        kp_detector,
        relative=True,
        adapt_movement_scale=True,
        cpu=use_cpu
    )

    output_path = 'output.mp4'
    imageio.mimsave(output_path, [img_as_ubyte(f) for f in predictions], fps=fps)
    return output_path

with gr.Blocks() as demo:
    gr.Markdown("First Order Motion Model")
    gr.Markdown("Anima una foto de cara con el movimiento de un vídeo conductor.")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="numpy", label="Imagen fuente (foto de cara)")
            input_video = gr.Video(label="Vídeo conductor")
            btn = gr.Button("Generar animación", variant="primary")
        with gr.Column():
            output_video = gr.Video(label="Resultado")

    btn.click(fn=generate, inputs=[input_image, input_video], outputs=output_video)

    gr.Examples(
        examples=[
            ["LinusTechTips.jpg", "DameDane.mp4"]
        ],
        inputs=[input_image, input_video],
        label="Ejemplo"
    )

demo.launch()