apol/ArcaneTest

1"""2Thanks to nateraw for making this scape happen! 3This code has been mostly taken from https://huggingface.co/spaces/nateraw/animegan-v2-for-videos/tree/main4"""5import os6os.system("wget https://github.com/Sxela/ArcaneGAN/releases/download/v0.3/ArcaneGANv0.3.jit")7 8import sys9from subprocess import call10def run_cmd(command):11    try:12        print(command)13        call(command, shell=True)14    except KeyboardInterrupt:15        print("Process interrupted")16        sys.exit(1)17 18print("⬇️ Installing latest gradio==2.4.7b9")19run_cmd("pip install --upgrade pip")20run_cmd('pip install gradio==2.4.7b9')21 22import gc23import math24 25 26import gradio as gr27import numpy as np28import torch29from encoded_video import EncodedVideo, write_video30from PIL import Image31from torchvision.transforms.functional import center_crop, to_tensor32 33 34 35 36print("🧠 Loading Model...")37model = torch.jit.load('./ArcaneGANv0.3.jit').cuda().eval().half()38 39# This function is taken from pytorchvideo!40def uniform_temporal_subsample(x: torch.Tensor, num_samples: int, temporal_dim: int = -3) -> torch.Tensor:41    """42    Uniformly subsamples num_samples indices from the temporal dimension of the video.43    When num_samples is larger than the size of temporal dimension of the video, it44    will sample frames based on nearest neighbor interpolation.45    Args:46        x (torch.Tensor): A video tensor with dimension larger than one with torch47            tensor type includes int, long, float, complex, etc.48        num_samples (int): The number of equispaced samples to be selected49        temporal_dim (int): dimension of temporal to perform temporal subsample.50    Returns:51        An x-like Tensor with subsampled temporal dimension.52    """53    t = x.shape[temporal_dim]54    assert num_samples > 0 and t > 055    # Sample by nearest neighbor interpolation if num_samples > t.56    indices = torch.linspace(0, t - 1, num_samples)57    indices = torch.clamp(indices, 0, t - 1).long()58    return torch.index_select(x, temporal_dim, indices)59 60 61# This function is taken from pytorchvideo!62def short_side_scale(63    x: torch.Tensor,64    size: int,65    interpolation: str = "bilinear",66) -> torch.Tensor:67    """68    Determines the shorter spatial dim of the video (i.e. width or height) and scales69    it to the given size. To maintain aspect ratio, the longer side is then scaled70    accordingly.71    Args:72        x (torch.Tensor): A video tensor of shape (C, T, H, W) and type torch.float32.73        size (int): The size the shorter side is scaled to.74        interpolation (str): Algorithm used for upsampling,75            options: nearest' | 'linear' | 'bilinear' | 'bicubic' | 'trilinear' | 'area'76    Returns:77        An x-like Tensor with scaled spatial dims.78    """79    assert len(x.shape) == 480    assert x.dtype == torch.float3281    c, t, h, w = x.shape82    if w < h:83        new_h = int(math.floor((float(h) / w) * size))84        new_w = size85    else:86        new_h = size87        new_w = int(math.floor((float(w) / h) * size))88 89    return torch.nn.functional.interpolate(x, size=(new_h, new_w), mode=interpolation, align_corners=False)90 91means = [0.485, 0.456, 0.406]92stds = [0.229, 0.224, 0.225]93 94from torchvision import transforms95norm = transforms.Normalize(means,stds)96 97norms = torch.tensor(means)[None,:,None,None].cuda()98stds = torch.tensor(stds)[None,:,None,None].cuda()99 100def inference_step(vid, start_sec, duration, out_fps, interpolate):101    clip = vid.get_clip(start_sec, start_sec + duration)102    video_arr = torch.from_numpy(clip['video']).permute(3, 0, 1, 2)103    audio_arr = np.expand_dims(clip['audio'], 0)104    audio_fps = None if not vid._has_audio else vid._container.streams.audio[0].sample_rate105 106    x = uniform_temporal_subsample(video_arr,  duration * out_fps)107    x = center_crop(short_side_scale(x, 512), 512)108    x /= 255.109    x = x.permute(1, 0, 2, 3)110    x = norm(x)111 112    with torch.no_grad():113        output = model(x.to('cuda').half())114        output = (output * stds + norms).clip(0, 1) * 255.115 116        output_video = output.permute(0, 2, 3, 1).float().detach().cpu().numpy()117        if interpolate == 'Yes': output_video[1:] = output_video[1:]*(0.5) + output_video[:-1]*(0.5)118    119    return output_video, audio_arr, out_fps, audio_fps120 121 122def predict_fn(filepath, start_sec, duration, out_fps, interpolate):123    # out_fps=12124    vid = EncodedVideo.from_path(filepath)125    for i in range(duration):126        video, audio, fps, audio_fps = inference_step(127            vid = vid,128            start_sec = i + start_sec,129            duration = 1,130            out_fps = out_fps,131            interpolate = interpolate132        )133        gc.collect()134        if i == 0:135            video_all = video136            audio_all = audio137        else:138            video_all = np.concatenate((video_all, video))139            audio_all = np.hstack((audio_all, audio))140 141    write_video(142        'out.mp4',143        video_all,144        fps=fps,145        audio_array=audio_all,146        audio_fps=audio_fps,147        audio_codec='aac'148    )149 150    del video_all151    del audio_all152    153    return 'out.mp4'154 155 156title = "ArcaneGAN"157description = "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."158article = "<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>"159 160 161gr.Interface(162    predict_fn,163    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')],164    outputs=gr.outputs.Video(),165    title='ArcaneGAN On Videos',166    description = description,167    article = article,168    enable_queue=True,169    examples=[170        ['obama.webm', 23, 10, 30],171    ],172    allow_flagging=False173).launch()174