VideoCrafterXtend / predict.py
ychenhq's picture
Upload 5 files
069147a verified
raw
history blame
5.57 kB
# Prediction interface for Cog ⚙️
# https://github.com/replicate/cog/blob/main/docs/python.md
import os
import sys
import argparse
import random
from omegaconf import OmegaConf
from einops import rearrange, repeat
import torch
import torchvision
from pytorch_lightning import seed_everything
from cog import BasePredictor, Input, Path
sys.path.insert(0, "scripts/evaluation")
from funcs import (
batch_ddim_sampling,
load_model_checkpoint,
load_image_batch,
get_filelist,
)
from utils.utils import instantiate_from_config
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory to make running multiple predictions efficient"""
ckpt_path_base = "checkpoints/base_1024_v1/model.ckpt"
config_base = "configs/inference_t2v_1024_v1.0.yaml"
ckpt_path_i2v = "checkpoints/i2v_512_v1/model.ckpt"
config_i2v = "configs/inference_i2v_512_v1.0.yaml"
config_base = OmegaConf.load(config_base)
model_config_base = config_base.pop("model", OmegaConf.create())
self.model_base = instantiate_from_config(model_config_base)
self.model_base = self.model_base.cuda()
self.model_base = load_model_checkpoint(self.model_base, ckpt_path_base)
self.model_base.eval()
config_i2v = OmegaConf.load(config_i2v)
model_config_i2v = config_i2v.pop("model", OmegaConf.create())
self.model_i2v = instantiate_from_config(model_config_i2v)
self.model_i2v = self.model_i2v.cuda()
self.model_i2v = load_model_checkpoint(self.model_i2v, ckpt_path_i2v)
self.model_i2v.eval()
def predict(
self,
task: str = Input(
description="Choose the task.",
choices=["text2video", "image2video"],
default="text2video",
),
prompt: str = Input(
description="Prompt for video generation.",
default="A tiger walks in the forest, photorealistic, 4k, high definition.",
),
image: Path = Input(
description="Input image for image2video task.", default=None
),
ddim_steps: int = Input(description="Number of denoising steps.", default=50),
unconditional_guidance_scale: float = Input(
description="Classifier-free guidance scale.", default=12.0
),
seed: int = Input(
description="Random seed. Leave blank to randomize the seed", default=None
),
save_fps: int = Input(
description="Frame per second for the generated video.", default=10
),
) -> Path:
width = 1024 if task == "text2video" else 512
height = 576 if task == "text2video" else 320
model = self.model_base if task == "text2video" else self.model_i2v
if task == "image2video":
assert image is not None, "Please provide image for image2video generation."
if seed is None:
seed = int.from_bytes(os.urandom(2), "big")
print(f"Using seed: {seed}")
seed_everything(seed)
args = argparse.Namespace(
mode="base" if task == "text2video" else "i2v",
savefps=save_fps,
n_samples=1,
ddim_steps=ddim_steps,
ddim_eta=1.0,
bs=1,
height=height,
width=width,
frames=-1,
fps=28 if task == "text2video" else 8,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_guidance_scale_temporal=None,
)
## latent noise shape
h, w = args.height // 8, args.width // 8
frames = model.temporal_length if args.frames < 0 else args.frames
channels = model.channels
batch_size = 1
noise_shape = [batch_size, channels, frames, h, w]
fps = torch.tensor([args.fps] * batch_size).to(model.device).long()
prompts = [prompt]
text_emb = model.get_learned_conditioning(prompts)
if args.mode == "base":
cond = {"c_crossattn": [text_emb], "fps": fps}
elif args.mode == "i2v":
cond_images = load_image_batch([str(image)], (args.height, args.width))
cond_images = cond_images.to(model.device)
img_emb = model.get_image_embeds(cond_images)
imtext_cond = torch.cat([text_emb, img_emb], dim=1)
cond = {"c_crossattn": [imtext_cond], "fps": fps}
else:
raise NotImplementedError
## inference
batch_samples = batch_ddim_sampling(
model,
cond,
noise_shape,
args.n_samples,
args.ddim_steps,
args.ddim_eta,
args.unconditional_guidance_scale,
)
out_path = "/tmp/output.mp4"
vid_tensor = batch_samples[0]
video = vid_tensor.detach().cpu()
video = torch.clamp(video.float(), -1.0, 1.0)
video = video.permute(2, 0, 1, 3, 4) # t,n,c,h,w
frame_grids = [
torchvision.utils.make_grid(framesheet, nrow=int(args.n_samples))
for framesheet in video
] # [3, 1*h, n*w]
grid = torch.stack(frame_grids, dim=0) # stack in temporal dim [t, 3, n*h, w]
grid = (grid + 1.0) / 2.0
grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1)
torchvision.io.write_video(
out_path,
grid,
fps=args.savefps,
video_codec="h264",
options={"crf": "10"},
)
return Path(out_path)