Spaces:
Running
Running
| import torch | |
| from diffusers.image_processor import VaeImageProcessor | |
| from torch.nn import functional as F | |
| import cv2 | |
| import utils | |
| from rife.pytorch_msssim import ssim_matlab | |
| import numpy as np | |
| import logging | |
| import skvideo.io | |
| from rife.RIFE_HDv3 import Model | |
| logger = logging.getLogger(__name__) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| def pad_image(img, scale): | |
| _, _, h, w = img.shape | |
| tmp = max(32, int(32 / scale)) | |
| ph = ((h - 1) // tmp + 1) * tmp | |
| pw = ((w - 1) // tmp + 1) * tmp | |
| padding = (0, 0, pw - w, ph - h) | |
| return F.pad(img, padding) | |
| def make_inference(model, I0, I1, upscale_amount, n): | |
| middle = model.inference(I0, I1, upscale_amount) | |
| if n == 1: | |
| return [middle] | |
| first_half = make_inference(model, I0, middle, upscale_amount, n=n // 2) | |
| second_half = make_inference(model, middle, I1, upscale_amount, n=n // 2) | |
| if n % 2: | |
| return [*first_half, middle, *second_half] | |
| else: | |
| return [*first_half, *second_half] | |
| def ssim_interpolation_rife(model, samples, exp=1, upscale_amount=1, output_device="cpu"): | |
| print(f"samples dtype:{samples.dtype}") | |
| print(f"samples shape:{samples.shape}") | |
| output = [] | |
| # [f, c, h, w] | |
| for b in range(samples.shape[0]): | |
| frame = samples[b : b + 1] | |
| _, _, h, w = frame.shape | |
| I0 = samples[b : b + 1] | |
| I1 = samples[b + 1 : b + 2] if b + 2 < samples.shape[0] else samples[-1:] | |
| I1 = pad_image(I1, upscale_amount) | |
| # [c, h, w] | |
| I0_small = F.interpolate(I0, (32, 32), mode="bilinear", align_corners=False) | |
| I1_small = F.interpolate(I1, (32, 32), mode="bilinear", align_corners=False) | |
| ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) | |
| if ssim > 0.996: | |
| I1 = I0 | |
| I1 = pad_image(I1, upscale_amount) | |
| I1 = make_inference(model, I0, I1, upscale_amount, 1) | |
| I1_small = F.interpolate(I1[0], (32, 32), mode="bilinear", align_corners=False) | |
| ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) | |
| frame = I1[0] | |
| I1 = I1[0] | |
| tmp_output = [] | |
| if ssim < 0.2: | |
| for i in range((2**exp) - 1): | |
| tmp_output.append(I0) | |
| else: | |
| tmp_output = make_inference(model, I0, I1, upscale_amount, 2**exp - 1) if exp else [] | |
| frame = pad_image(frame, upscale_amount) | |
| tmp_output = [frame] + tmp_output | |
| for i, frame in enumerate(tmp_output): | |
| output.append(frame.to(output_device)) | |
| return output | |
| def load_rife_model(model_path): | |
| torch.set_grad_enabled(False) | |
| if torch.cuda.is_available(): | |
| torch.backends.cudnn.enabled = True | |
| torch.backends.cudnn.benchmark = True | |
| torch.set_default_tensor_type(torch.cuda.FloatTensor) | |
| model = Model() | |
| model.load_model(model_path, -1) | |
| model.eval() | |
| print("Loaded v3.x HD model.") | |
| return model | |
| # Create a generator that yields each frame, similar to cv2.VideoCapture | |
| def frame_generator(video_capture): | |
| while True: | |
| ret, frame = video_capture.read() | |
| if not ret: | |
| break | |
| yield frame | |
| video_capture.release() | |
| def rife_inference_with_path(model, video_path): | |
| video_capture = cv2.VideoCapture(video_path) | |
| tot_frame = video_capture.get(cv2.CAP_PROP_FRAME_COUNT) | |
| pt_frame_data = [] | |
| pt_frame = skvideo.io.vreader(video_path) | |
| for frame in pt_frame: | |
| pt_frame_data.append( | |
| torch.from_numpy(np.transpose(frame, (2, 0, 1))).to("cpu", non_blocking=True).float() / 255.0 | |
| ) | |
| pt_frame = torch.from_numpy(np.stack(pt_frame_data)) | |
| pt_frame = pt_frame.to(device) | |
| pbar = utils.ProgressBar(tot_frame, desc="RIFE inference") | |
| frames = ssim_interpolation_rife(model, pt_frame) | |
| pt_image = torch.stack([frames[i].squeeze(0) for i in range(len(frames))]) | |
| image_np = VaeImageProcessor.pt_to_numpy(pt_image) # (to [49, 512, 480, 3]) | |
| image_pil = VaeImageProcessor.numpy_to_pil(image_np) | |
| video_path = utils.save_video(image_pil, fps=16) | |
| if pbar: | |
| pbar.update(1) | |
| return video_path | |
| def rife_inference_with_latents(model, latents): | |
| pbar = utils.ProgressBar(latents.shape[1], desc="RIFE inference") | |
| rife_results = [] | |
| latents = latents.to(device) | |
| for i in range(latents.size(0)): | |
| # [f, c, w, h] | |
| latent = latents[i] | |
| frames = ssim_interpolation_rife(model, latent) | |
| pt_image = torch.stack([frames[i].squeeze(0) for i in range(len(frames))]) # (to [f, c, w, h]) | |
| rife_results.append(pt_image) | |
| return torch.stack(rife_results) | |