104 lines
4.2 KiB
Python
104 lines
4.2 KiB
Python
# General
|
|
import os
|
|
from os.path import join as opj
|
|
import argparse
|
|
import datetime
|
|
from pathlib import Path
|
|
import torch
|
|
import gradio as gr
|
|
import tempfile
|
|
import yaml
|
|
from t2v_enhanced.model.video_ldm import VideoLDM
|
|
from typing import List, Optional
|
|
from model.callbacks import SaveConfigCallback
|
|
from PIL.Image import Image, fromarray
|
|
|
|
from einops import rearrange, repeat
|
|
|
|
import sys
|
|
sys.path.append("thirdparty")
|
|
from modelscope.pipelines import pipeline
|
|
from modelscope.outputs import OutputKeys
|
|
import imageio
|
|
import pathlib
|
|
import numpy as np
|
|
|
|
# Utilities
|
|
from inference_utils import *
|
|
from model_init import *
|
|
from model_func import *
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--prompt', type=str, default="A cat running on the street", help="The prompt to guide video generation.")
|
|
parser.add_argument('--image', type=str, default="", help="Path to image conditioning.")
|
|
# parser.add_argument('--video', type=str, default="", help="Path to video conditioning.")
|
|
parser.add_argument('--base_model', type=str, default="ModelscopeT2V", help="Base model to generate first chunk from", choices=["ModelscopeT2V", "AnimateDiff", "SVD"])
|
|
parser.add_argument('--num_frames', type=int, default=24, help="The number of video frames to generate.")
|
|
parser.add_argument('--negative_prompt', type=str, default="", help="The prompt to guide what to not include in video generation.")
|
|
parser.add_argument('--num_steps', type=int, default=50, help="The number of denoising steps.")
|
|
parser.add_argument('--image_guidance', type=float, default=9.0, help="The guidance scale.")
|
|
|
|
parser.add_argument('--output_dir', type=str, default="results", help="Path where to save the generated videos.")
|
|
parser.add_argument('--device', type=str, default="cuda")
|
|
parser.add_argument('--seed', type=int, default=33, help="Random seed")
|
|
|
|
parser.add_argument('--chunk', type=int, default=56, help="chunk_size for randomized blending")
|
|
parser.add_argument('--overlap', type=int, default=32, help="overlap_size for randomized blending")
|
|
args = parser.parse_args()
|
|
|
|
|
|
Path(args.output_dir).mkdir(parents=True, exist_ok=True)
|
|
result_fol = Path(args.output_dir).absolute()
|
|
device = args.device
|
|
|
|
|
|
# --------------------------
|
|
# ----- Configurations -----
|
|
# --------------------------
|
|
ckpt_file_streaming_t2v = Path("checkpoints/streaming_t2v.ckpt").absolute()
|
|
cfg_v2v = {'downscale': 1, 'upscale_size': (1280, 720), 'model_id': 'damo/Video-to-Video', 'pad': True}
|
|
|
|
# --------------------------
|
|
# ----- Initialization -----
|
|
# --------------------------
|
|
stream_cli, stream_model = init_streamingt2v_model(ckpt_file_streaming_t2v, result_fol)
|
|
if args.base_model == "ModelscopeT2V":
|
|
model = init_modelscope(device)
|
|
elif args.base_model == "AnimateDiff":
|
|
model = init_animatediff(device)
|
|
elif args.base_model == "SVD":
|
|
model = init_svd(device)
|
|
sdxl_model = init_sdxl(device)
|
|
|
|
msxl_model = init_v2v_model(cfg_v2v)
|
|
|
|
inference_generator = torch.Generator(device="cuda")
|
|
|
|
|
|
# ------------------
|
|
# ----- Inputs -----
|
|
# ------------------
|
|
now = datetime.datetime.now()
|
|
name = args.prompt[:100].replace(" ", "_") + "_" + str(now.time()).replace(":", "_").replace(".", "_")
|
|
|
|
inference_generator = torch.Generator(device="cuda")
|
|
inference_generator.manual_seed(args.seed)
|
|
|
|
if args.base_model == "ModelscopeT2V":
|
|
short_video = ms_short_gen(args.prompt, model, inference_generator)
|
|
elif args.base_model == "AnimateDiff":
|
|
short_video = ad_short_gen(args.prompt, model, inference_generator)
|
|
elif args.base_model == "SVD":
|
|
short_video = svd_short_gen(args.image, args.prompt, model, sdxl_model, inference_generator)
|
|
|
|
n_autoreg_gen = (args.num_frames-8)//8
|
|
stream_long_gen(args.prompt, short_video, n_autoreg_gen, args.negative_prompt, args.seed, args.num_steps, args.image_guidance, name, stream_cli, stream_model)
|
|
|
|
|
|
if args.num_frames > 80:
|
|
video2video_randomized(args.prompt, opj(result_fol, name+".mp4"), result_fol, cfg_v2v, msxl_model, chunk_size=args.chunk, overlap_size=args.overlap)
|
|
else:
|
|
video2video(args.prompt, opj(result_fol, name+".mp4"), result_fol, cfg_v2v, msxl_model)
|