File size: 6,206 Bytes
8810a39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e618b8f
8810a39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5b4dfa2
8810a39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5b4dfa2
8810a39
 
 
 
5b4dfa2
8810a39
 
 
 
 
 
 
5b4dfa2
 
8810a39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
defabb6
8810a39
 
 
 
 
9371865
8810a39
 
4326472
8810a39
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
"""
Thanks to nateraw for making this scape happen! 
This code has been mostly taken from https://huggingface.co/spaces/nateraw/animegan-v2-for-videos/tree/main
"""
import os
os.system("wget https://github.com/Sxela/ArcaneGAN/releases/download/v0.3/ArcaneGANv0.3.jit")

import sys
from subprocess import call
def run_cmd(command):
    try:
        print(command)
        call(command, shell=True)
    except KeyboardInterrupt:
        print("Process interrupted")
        sys.exit(1)

print("⬇️ Installing latest gradio==2.4.7b9")
run_cmd("pip install --upgrade pip")
run_cmd('pip install gradio==2.4.7b9')

import gc
import math


import gradio as gr
import numpy as np
import torch
from encoded_video import EncodedVideo, write_video
from PIL import Image
from torchvision.transforms.functional import center_crop, to_tensor




print("🧠 Loading Model...")
model = torch.jit.load('./ArcaneGANv0.3.jit').cuda().eval().half()

# This function is taken from pytorchvideo!
def uniform_temporal_subsample(x: torch.Tensor, num_samples: int, temporal_dim: int = -3) -> torch.Tensor:
    """
    Uniformly subsamples num_samples indices from the temporal dimension of the video.
    When num_samples is larger than the size of temporal dimension of the video, it
    will sample frames based on nearest neighbor interpolation.
    Args:
        x (torch.Tensor): A video tensor with dimension larger than one with torch
            tensor type includes int, long, float, complex, etc.
        num_samples (int): The number of equispaced samples to be selected
        temporal_dim (int): dimension of temporal to perform temporal subsample.
    Returns:
        An x-like Tensor with subsampled temporal dimension.
    """
    t = x.shape[temporal_dim]
    assert num_samples > 0 and t > 0
    # Sample by nearest neighbor interpolation if num_samples > t.
    indices = torch.linspace(0, t - 1, num_samples)
    indices = torch.clamp(indices, 0, t - 1).long()
    return torch.index_select(x, temporal_dim, indices)


# This function is taken from pytorchvideo!
def short_side_scale(
    x: torch.Tensor,
    size: int,
    interpolation: str = "bilinear",
) -> torch.Tensor:
    """
    Determines the shorter spatial dim of the video (i.e. width or height) and scales
    it to the given size. To maintain aspect ratio, the longer side is then scaled
    accordingly.
    Args:
        x (torch.Tensor): A video tensor of shape (C, T, H, W) and type torch.float32.
        size (int): The size the shorter side is scaled to.
        interpolation (str): Algorithm used for upsampling,
            options: nearest' | 'linear' | 'bilinear' | 'bicubic' | 'trilinear' | 'area'
    Returns:
        An x-like Tensor with scaled spatial dims.
    """
    assert len(x.shape) == 4
    assert x.dtype == torch.float32
    c, t, h, w = x.shape
    if w < h:
        new_h = int(math.floor((float(h) / w) * size))
        new_w = size
    else:
        new_h = size
        new_w = int(math.floor((float(w) / h) * size))

    return torch.nn.functional.interpolate(x, size=(new_h, new_w), mode=interpolation, align_corners=False)

means = [0.485, 0.456, 0.406]
stds = [0.229, 0.224, 0.225]

from torchvision import transforms
norm = transforms.Normalize(means,stds)

norms = torch.tensor(means)[None,:,None,None].cuda()
stds = torch.tensor(stds)[None,:,None,None].cuda()

def inference_step(vid, start_sec, duration, out_fps, interpolate):
    clip = vid.get_clip(start_sec, start_sec + duration)
    video_arr = torch.from_numpy(clip['video']).permute(3, 0, 1, 2)
    audio_arr = np.expand_dims(clip['audio'], 0)
    audio_fps = None if not vid._has_audio else vid._container.streams.audio[0].sample_rate

    x = uniform_temporal_subsample(video_arr,  duration * out_fps)
    x = center_crop(short_side_scale(x, 512), 512)
    x /= 255.
    x = x.permute(1, 0, 2, 3)
    x = norm(x)

    with torch.no_grad():
        output = model(x.to('cuda').half())
        output = (output * stds + norms).clip(0, 1) * 255.

        output_video = output.permute(0, 2, 3, 1).float().detach().cpu().numpy()
        if interpolate == 'Yes': output_video[1:] = output_video[1:]*(0.5) + output_video[:-1]*(0.5)
    
    return output_video, audio_arr, out_fps, audio_fps


def predict_fn(filepath, start_sec, duration, out_fps, interpolate):
    # out_fps=12
    vid = EncodedVideo.from_path(filepath)
    for i in range(duration):
        video, audio, fps, audio_fps = inference_step(
            vid = vid,
            start_sec = i + start_sec,
            duration = 1,
            out_fps = out_fps,
            interpolate = interpolate
        )
        gc.collect()
        if i == 0:
            video_all = video
            audio_all = audio
        else:
            video_all = np.concatenate((video_all, video))
            audio_all = np.hstack((audio_all, audio))

    write_video(
        'out.mp4',
        video_all,
        fps=fps,
        audio_array=audio_all,
        audio_fps=audio_fps,
        audio_codec='aac'
    )

    del video_all
    del audio_all
    
    return 'out.mp4'


title = "ArcaneGAN"
description = "Gradio demo for ArcaneGAN, video to Arcane style. To use it, simply upload your video, or click on an example below. Follow me on twitter for more info and updates."
article = "<div style='text-align: center;'>ArcaneGan by <a href='https://twitter.com/devdef' target='_blank'>Alex Spirin</a> | <a href='https://github.com/Sxela/ArcaneGAN' target='_blank'>Github Repo</a> | <center><img src='https://visitor-badge.glitch.me/badge?page_id=sxela_arcanegan_video_hf' alt='visitor badge'></center></div>"


gr.Interface(
    predict_fn,
    inputs=[gr.inputs.Video(), gr.inputs.Slider(minimum=0, maximum=300, step=1, default=0), gr.inputs.Slider(minimum=1, maximum=10, step=1, default=2), gr.inputs.Slider(minimum=12, maximum=30, step=6, default=24), gr.inputs.Radio(choices=['Yes','No'], type="value", default='Yes', label='Remove flickering')],
    outputs=gr.outputs.Video(),
    title='ArcaneGAN On Videos',
    description = description,
    article = article,
    enable_queue=True,
    examples=[
        ['obama.webm', 23, 10, 30],
    ],
    allow_flagging=False
).launch()