Spaces:
Running
Running
import random | |
import gradio as gr | |
import huggingface_hub | |
import imageio | |
import numpy as np | |
import onnxruntime as rt | |
from numpy.random import RandomState | |
from skimage import transform | |
class Model: | |
def __init__(self): | |
self.g_synthesis = None | |
self.g_mapping = None | |
self.load_models() | |
def load_models(self): | |
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] | |
g_mapping_path = huggingface_hub.hf_hub_download("skytnt/waifu-gan", "g_mapping.onnx") | |
g_synthesis_path = huggingface_hub.hf_hub_download("skytnt/waifu-gan", "g_synthesis.onnx") | |
self.g_mapping = rt.InferenceSession(g_mapping_path, providers=providers) | |
self.g_synthesis = rt.InferenceSession(g_synthesis_path, providers=providers) | |
def get_img(self, w): | |
img = self.g_synthesis.run(None, {'w': w})[0] | |
return (img.transpose(0, 2, 3, 1) * 127.5 + 128).clip(0, 255).astype(np.uint8)[0] | |
def get_w(self, z, psi1, psi2): | |
return self.g_mapping.run(None, {'z': z, 'psi': np.asarray([psi1, psi2], dtype=np.float32)})[0] | |
def gen_video(self, w1, w2, path, frame_num=10): | |
video = imageio.get_writer(path, mode='I', fps=frame_num // 2, codec='libx264', bitrate='16M') | |
lin = np.linspace(0, 1, frame_num) | |
for i in range(0, frame_num): | |
img = self.get_img(((1 - lin[i]) * w1) + (lin[i] * w2)) | |
video.append_data(img) | |
video.close() | |
def get_thumbnail(img): | |
img_new = np.full((192, 288, 3), 200, dtype=np.uint8) | |
img_new[:, 80:208] = transform.resize(img, (192, 128), preserve_range=True) | |
return img_new | |
def gen_fn(method, seed, psi1, psi2): | |
if method == 0: | |
seed = random.randint(0, 2 ** 32 -1) | |
z = RandomState(int(seed)).randn(1, 1024) | |
w = model.get_w(z.astype(dtype=np.float32), psi1, psi2) | |
img_out = model.get_img(w) | |
return img_out, seed, w, get_thumbnail(img_out) | |
def gen_video_fn(w1, w2, frame): | |
if w1 is None or w2 is None: | |
return None | |
model.gen_video(w1, w2, "video.mp4", int(frame)) | |
return "video.mp4" | |
if __name__ == '__main__': | |
model = Model() | |
app = gr.Blocks() | |
with app: | |
gr.Markdown("# Waifu GAN\n\n" | |
"![visitor badge](https://api.visitorbadge.io/api/visitors?path=skytnt.waifu-gan&countColor=%23263759&style=flat&labelStyle=lower)\n\n") | |
with gr.Tabs(): | |
with gr.TabItem("generate image"): | |
with gr.Row(): | |
with gr.Column(): | |
with gr.Row(): | |
gen_input1 = gr.Radio(label="method", value="random", | |
choices=["random", "use seed"], type="index") | |
gen_input2 = gr.Slider(minimum=0, maximum=2 ** 32 - 1, step=1, value=0, | |
label="seed") | |
gen_input3 = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, label="truncation psi 1") | |
gen_input4 = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, label="truncation psi 2") | |
with gr.Group(): | |
gen_submit = gr.Button("Generate", variant="primary") | |
with gr.Column(): | |
gen_output1 = gr.Image(label="output image") | |
select_img_input_w1 = gr.State() | |
select_img_input_img1 = gr.State() | |
with gr.TabItem("generate video"): | |
with gr.Row(): | |
with gr.Column(): | |
gr.Markdown("## generate video between 2 images") | |
with gr.Row(): | |
with gr.Column(): | |
gr.Markdown("please select image 1") | |
select_img1_dropdown = gr.Radio(label="source", value="current generated image", | |
choices=["current generated image"], type="index") | |
with gr.Group(): | |
select_img1_button = gr.Button("Select", variant="primary") | |
select_img1_output_img = gr.Image(label="selected image 1") | |
select_img1_output_w = gr.State() | |
with gr.Column(): | |
gr.Markdown("please select image 2") | |
select_img2_dropdown = gr.Radio(label="source", value="current generated image", | |
choices=["current generated image"], type="index") | |
with gr.Group(): | |
select_img2_button = gr.Button("Select", variant="primary") | |
select_img2_output_img = gr.Image(label="selected image 2") | |
select_img2_output_w = gr.State() | |
generate_video_frame = gr.Slider(minimum=10, maximum=30, step=1, label="frame", value=15) | |
with gr.Group(): | |
generate_video_button = gr.Button("Generate", variant="primary") | |
with gr.Column(): | |
generate_video_output = gr.Video(label="output video") | |
gen_submit.click(gen_fn, [gen_input1, gen_input2, gen_input3, gen_input4], | |
[gen_output1, gen_input2, select_img_input_w1, select_img_input_img1]) | |
select_img1_button.click(lambda i, img1, w1: (img1, w1), | |
[select_img1_dropdown, select_img_input_img1, select_img_input_w1], | |
[select_img1_output_img, select_img1_output_w]) | |
select_img2_button.click(lambda i, img1, w1: (img1, w1), | |
[select_img2_dropdown, select_img_input_img1, select_img_input_w1], | |
[select_img2_output_img, select_img2_output_w]) | |
generate_video_button.click(gen_video_fn, | |
[select_img1_output_w, select_img2_output_w, generate_video_frame], | |
[generate_video_output]) | |
app.launch() | |