Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,105 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import tempfile
|
2 |
+
|
3 |
+
import cv2
|
4 |
+
import ffmpegio
|
5 |
+
import gradio as gr
|
6 |
+
import numpy as np
|
7 |
+
import omegaconf
|
8 |
+
import tensorflow as tf
|
9 |
+
from pyprojroot.pyprojroot import here
|
10 |
+
from huggingface_hub import hf_hub_url, cached_download
|
11 |
+
|
12 |
+
from ganime.model.vqgan_clean.experimental.net2net_v3 import Net2Net
|
13 |
+
|
14 |
+
IMAGE_SHAPE = (64, 128, 3)
|
15 |
+
|
16 |
+
vqgan_path = cached_download(
|
17 |
+
hf_hub_url("Kurokabe", "VQGAN_Kimetsu-no-yaiba_Tensorflow/vqgan_kny_image_full")
|
18 |
+
)
|
19 |
+
gpt_path = cached_download(
|
20 |
+
hf_hub_url("Kurokabe", "GANime_Kimetsu-no-yaiba_Tensorflow/ganime_kny_video_full")
|
21 |
+
)
|
22 |
+
|
23 |
+
cfg = omegaconf.OmegaConf.load(here("configs/kny_video_gpt2_large_gradio.yaml"))
|
24 |
+
cfg["model"]["first_stage_config"]["checkpoint_path"] = vqgan_path + "/checkpoint"
|
25 |
+
cfg["model"]["transformer_config"]["checkpoint_path"] = gpt_path + "/checkpoint"
|
26 |
+
|
27 |
+
model = Net2Net(**cfg["model"], trainer_config=cfg["train"], num_replicas=1)
|
28 |
+
model.first_stage_model.build((20, *IMAGE_SHAPE))
|
29 |
+
|
30 |
+
|
31 |
+
# def save_video(video):
|
32 |
+
# b, f, h, w, c = 1, 20, 500, 500, 3
|
33 |
+
|
34 |
+
# # filename = output_file.name
|
35 |
+
# filename = "./test_video.mp4"
|
36 |
+
# images = []
|
37 |
+
# for i in range(f):
|
38 |
+
# # image = video[0][i].numpy()
|
39 |
+
# # image = 255 * image # Now scale by 255
|
40 |
+
# # image = image.astype(np.uint8)
|
41 |
+
# images.append(np.random.randint(0, 255, (h, w, c), dtype=np.uint8))
|
42 |
+
|
43 |
+
# ffmpegio.video.write(filename, 20, np.array(images), overwrite=True)
|
44 |
+
# return filename
|
45 |
+
|
46 |
+
|
47 |
+
def save_video(video):
|
48 |
+
output_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
|
49 |
+
b, f, h, w, c = video.shape
|
50 |
+
|
51 |
+
filename = output_file.name
|
52 |
+
video = video.numpy()
|
53 |
+
video = video * 255
|
54 |
+
video = video.astype(np.uint8)
|
55 |
+
ffmpegio.video.write(filename, 20, video, overwrite=True)
|
56 |
+
return filename
|
57 |
+
|
58 |
+
|
59 |
+
def resize_if_necessary(image):
|
60 |
+
if image.shape[0] != 64 and image.shape[1] != 128:
|
61 |
+
image = tf.image.resize(image, (64, 128))
|
62 |
+
return image
|
63 |
+
|
64 |
+
|
65 |
+
def normalize(image):
|
66 |
+
image = (tf.cast(image, tf.float32) / 127.5) - 1
|
67 |
+
|
68 |
+
return image
|
69 |
+
|
70 |
+
|
71 |
+
def generate(first, last, n_frames):
|
72 |
+
# n_frames = 20
|
73 |
+
n_frames = int(n_frames)
|
74 |
+
first = resize_if_necessary(first)
|
75 |
+
last = resize_if_necessary(last)
|
76 |
+
first = normalize(first)
|
77 |
+
last = normalize(last)
|
78 |
+
data = {
|
79 |
+
"first_frame": np.expand_dims(first, axis=0),
|
80 |
+
"last_frame": np.expand_dims(last, axis=0),
|
81 |
+
"y": None,
|
82 |
+
"n_frames": [n_frames],
|
83 |
+
"remaining_frames": [list(reversed(range(n_frames)))],
|
84 |
+
}
|
85 |
+
generated = model.predict(data)
|
86 |
+
|
87 |
+
return save_video(generated)
|
88 |
+
|
89 |
+
|
90 |
+
gr.Interface(
|
91 |
+
generate,
|
92 |
+
inputs=[
|
93 |
+
gr.Image(label="Upload the first image"),
|
94 |
+
gr.Image(label="Upload the last image"),
|
95 |
+
gr.Slider(
|
96 |
+
label="Number of frame to generate",
|
97 |
+
minimum=15,
|
98 |
+
maximum=100,
|
99 |
+
value=15,
|
100 |
+
step=1,
|
101 |
+
),
|
102 |
+
],
|
103 |
+
outputs="video",
|
104 |
+
title="Generate a video from the first and last frame",
|
105 |
+
).launch(share=True)
|