Spaces:
Running
Running
File size: 5,558 Bytes
566acd1 112171d 566acd1 112171d 566acd1 a90f4bc 566acd1 a90f4bc 566acd1 151f876 94539f1 566acd1 94539f1 566acd1 94539f1 151f876 |
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 |
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
device = "cuda" if torch.cuda.is_available() else "cpu"
print("🧠 Loading Model...")
model = torch.hub.load(
"AK391/animegan2-pytorch:main",
"generator",
pretrained=True,
device=device,
progress=True,
)
def face2paint(model: torch.nn.Module, img: Image.Image, size: int = 512, device: str = device):
w, h = img.size
s = min(w, h)
img = img.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
img = img.resize((size, size), Image.LANCZOS)
with torch.no_grad():
input = to_tensor(img).unsqueeze(0) * 2 - 1
output = model(input.to(device)).cpu()[0]
output = (output * 0.5 + 0.5).clip(0, 1) * 255.0
return output
# 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)
def inference_step(vid, start_sec, duration, out_fps):
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.0
x = x.permute(1, 0, 2, 3)
with torch.no_grad():
output = model(x.to(device)).detach().cpu()
output = (output * 0.5 + 0.5).clip(0, 1) * 255.0
output_video = output.permute(0, 2, 3, 1).numpy()
return output_video, audio_arr, out_fps, audio_fps
def predict_fn(filepath, start_sec, duration):
out_fps = 18
vid = EncodedVideo.from_path(filepath)
for i in range(duration):
print(f"🖼️ Processing step {i + 1}/{duration}...")
video, audio, fps, audio_fps = inference_step(vid=vid, start_sec=i + start_sec, duration=1, out_fps=out_fps)
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))
print(f"💾 Writing output video...")
try:
write_video('out.mp4', video_all, fps=fps, audio_array=audio_all, audio_fps=audio_fps, audio_codec='aac')
except:
print("❌ Error when writing with audio...trying without audio")
write_video('out.mp4', video_all, fps=fps)
print(f"✅ Done!")
del video_all
del audio_all
return 'out.mp4'
article = """
<p style='text-align: center'>
<a href='https://github.com/bryandlee/animegan2-pytorch' target='_blank'>Github Repo Pytorch</a>
</p>
"""
iface_file = gr.Interface(
predict_fn,
inputs=[
gr.Video(),
gr.Slider(minimum=0, maximum=300, step=1, value=0),
gr.Slider(minimum=1, maximum=10, step=1, value=2),
],
outputs=gr.Video(),
title='AnimeGANV2 On Videos',
description="Applying AnimeGAN-V2 to frames from video clips",
article=article,
examples=[
['driving.mp4', 0, 6],
['bella_poarch.mp4', 4, 8],
['obama.webm', 0, 4],
],
allow_flagging="never",
cache_examples="lazy",
delete_cache=(4000, 4000),
).queue(api_open=True).launch(show_error=True, show_api=True)
|