| import cv2 |
| import glob |
| import torch |
| import gradio as gr |
| import numpy as np |
| from huggingface_hub import hf_hub_download |
|
|
| from networks.amts import Model as AMTS |
| from networks.amtl import Model as AMTL |
| from networks.amtg import Model as AMTG |
| from utils import ( |
| img2tensor, tensor2img, |
| InputPadder, |
| check_dim_and_resize |
| ) |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| model_dict = { |
| 'AMT-S': AMTS, 'AMT-L': AMTL, 'AMT-G': AMTG |
| } |
|
|
| def img2vid(model_type, img0, img1, frame_ratio, iters): |
| model = model_dict[model_type]() |
| model.to(device) |
| ckpt_path = hf_hub_download(repo_id='lalala125/AMT', filename=f'{model_type.lower()}.pth') |
| ckpt = torch.load(ckpt_path, map_location=torch.device('cpu')) |
| model.load_state_dict(ckpt['state_dict']) |
| model.eval() |
| img0_t = img2tensor(img0).to(device) |
| img1_t = img2tensor(img1).to(device) |
| inputs = [img0_t, img1_t] |
| |
| |
| if device == 'cuda': |
| anchor_resolution = 1024 * 512 |
| anchor_memory = 1500 * 1024**2 |
| anchor_memory_bias = 2500 * 1024**2 |
| vram_avail = torch.cuda.get_device_properties(device).total_memory |
| else: |
| |
| anchor_resolution = 8192*8192 |
| anchor_memory = 1 |
| anchor_memory_bias = 0 |
| vram_avail = 1 |
| embt = torch.tensor(1/2).float().view(1, 1, 1, 1).to(device) |
|
|
| inputs = check_dim_and_resize(inputs) |
| h, w = inputs[0].shape[-2:] |
| scale = anchor_resolution / (h * w) * np.sqrt((vram_avail - anchor_memory_bias) / anchor_memory) |
| scale = 1 if scale > 1 else scale |
| scale = 1 / np.floor(1 / np.sqrt(scale) * 16) * 16 |
| if scale < 1: |
| print(f"Due to the limited VRAM, the video will be scaled by {scale:.2f}") |
| padding = int(16 / scale) |
| padder = InputPadder(inputs[0].shape, padding) |
| inputs = padder.pad(*inputs) |
|
|
| for i in range(iters): |
| print(f'Iter {i+1}. input_frames={len(inputs)} output_frames={2*len(inputs)-1}') |
| outputs = [inputs[0]] |
| for in_0, in_1 in zip(inputs[:-1], inputs[1:]): |
| in_0 = in_0.to(device) |
| in_1 = in_1.to(device) |
| with torch.no_grad(): |
| imgt_pred = model(in_0, in_1, embt, scale_factor=scale, eval=True)['imgt_pred'] |
| outputs += [imgt_pred.cpu(), in_1.cpu()] |
| inputs = outputs |
| outputs = padder.unpad(*outputs) |
| out_path = 'results' |
| size = outputs[0].shape[2:][::-1] |
| writer = cv2.VideoWriter(f'{out_path}/demo.mp4', cv2.VideoWriter_fourcc(*'mp4v'), frame_ratio, size) |
| for i, imgt_pred in enumerate(outputs): |
| imgt_pred = tensor2img(imgt_pred) |
| imgt_pred = cv2.cvtColor(imgt_pred, cv2.COLOR_RGB2BGR) |
| writer.write(imgt_pred) |
| writer.release() |
| return 'results/demo.mp4' |
|
|
| |
| def demo_img(): |
| with gr.Blocks() as demo: |
| with gr.Row(): |
| gr.Markdown('## Image Demo') |
| with gr.Row(): |
| gr.HTML( |
| """ |
| <div style="text-align: left; auto;"> |
| <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem"> |
| Description: With 2 input images, you can generate a short video from them. |
| </h3> |
| </div> |
| """) |
|
|
| with gr.Row(): |
| with gr.Column(): |
| img0 = gr.Image(label='Image0') |
| img1 = gr.Image(label='Image1') |
| with gr.Column(): |
| result = gr.Video(label="Generated Video") |
| with gr.Accordion('Advanced options', open=False): |
| ratio = gr.Slider(label='Multiple Ratio', |
| minimum=4, |
| maximum=7, |
| value=6, |
| step=1) |
| frame_ratio = gr.Slider(label='Frame Ratio', |
| minimum=8, |
| maximum=64, |
| value=16, |
| step=1) |
| model_type = gr.Radio(['AMT-S', 'AMT-L', 'AMT-G'], |
| label='Model Select', |
| value='AMT-S') |
| run_button = gr.Button(label='Run') |
| inputs = [ |
| model_type, |
| img0, |
| img1, |
| frame_ratio, |
| ratio, |
| ] |
|
|
| gr.Examples(examples=glob.glob("examples/*.png"), |
| inputs=img0, |
| label='Example images (drag them to input windows)', |
| run_on_click=False, |
| ) |
|
|
| run_button.click(fn=img2vid, |
| inputs=inputs, |
| outputs=result,) |
| return demo |
