236 lines
12 KiB
Python
236 lines
12 KiB
Python
import math
|
|
import os
|
|
import torch
|
|
import argparse
|
|
import torchvision
|
|
|
|
from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
|
|
EulerDiscreteScheduler, DPMSolverMultistepScheduler,
|
|
HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
|
|
DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
|
|
from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
|
|
from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
|
|
from omegaconf import OmegaConf
|
|
from torchvision.utils import save_image
|
|
from transformers import T5EncoderModel, MT5EncoderModel, UMT5EncoderModel, AutoTokenizer
|
|
|
|
|
|
from opensora.adaptor.modules import replace_with_fp32_forwards
|
|
from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
|
|
from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
|
|
# from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
|
|
|
|
from opensora.models.text_encoder import get_text_enc
|
|
from opensora.utils.utils import save_video_grid
|
|
|
|
from opensora.sample.pipeline_opensora import OpenSoraPipeline
|
|
|
|
import imageio
|
|
|
|
|
|
def main(args):
|
|
# torch.manual_seed(args.seed)
|
|
# torch.backends.cuda.matmul.allow_tf32 = False
|
|
weight_dtype = torch.bfloat16
|
|
device = torch.device(args.device)
|
|
|
|
# vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
|
|
vae = CausalVAEModelWrapper(args.ae_path)
|
|
vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
|
|
if args.enable_tiling:
|
|
vae.vae.enable_tiling()
|
|
vae.vae.tile_overlap_factor = args.tile_overlap_factor
|
|
vae.vae.tile_sample_min_size = 512
|
|
vae.vae.tile_latent_min_size = 64
|
|
vae.vae.tile_sample_min_size_t = 29
|
|
vae.vae.tile_latent_min_size_t = 8
|
|
if args.save_memory:
|
|
vae.vae.tile_sample_min_size = 256
|
|
vae.vae.tile_latent_min_size = 32
|
|
vae.vae.tile_sample_min_size_t = 29
|
|
vae.vae.tile_latent_min_size_t = 8
|
|
vae.vae_scale_factor = ae_stride_config[args.ae]
|
|
|
|
# if args.model_3d:
|
|
# transformer_model = OpenSoraT2V.from_pretrained(args.model_path, subfolder=args.version, cache_dir=args.cache_dir, low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
|
|
# else:
|
|
# transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, cache_dir=args.cache_dir, low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
|
|
|
|
if args.model_type == 'dit':
|
|
transformer_model = OpenSoraT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir,
|
|
low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
|
|
elif args.model_type == 'udit':
|
|
transformer_model = UDiTT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir, ignore_mismatched_sizes=True,
|
|
low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
|
|
else:
|
|
transformer_model = LatteT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir, low_cpu_mem_usage=False,
|
|
device_map=None, torch_dtype=weight_dtype)
|
|
|
|
text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir, low_cpu_mem_usage=True, torch_dtype=weight_dtype)
|
|
tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
|
|
# text_encoder = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir, low_cpu_mem_usage=True, torch_dtype=weight_dtype)
|
|
# tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
|
|
|
|
|
|
# set eval mode
|
|
transformer_model.eval()
|
|
vae.eval()
|
|
text_encoder.eval()
|
|
|
|
if args.sample_method == 'DDIM': #########
|
|
scheduler = DDIMScheduler(clip_sample=False)
|
|
elif args.sample_method == 'EulerDiscrete':
|
|
scheduler = EulerDiscreteScheduler()
|
|
elif args.sample_method == 'DDPM': #############
|
|
scheduler = DDPMScheduler(clip_sample=False)
|
|
elif args.sample_method == 'DPMSolverMultistep':
|
|
'''
|
|
DPM++ 2M DPMSolverMultistepScheduler
|
|
DPM++ 2M Karras DPMSolverMultistepScheduler init with use_karras_sigmas=True
|
|
DPM++ 2M SDE DPMSolverMultistepScheduler init with algorithm_type="sde-dpmsolver++"
|
|
DPM++ 2M SDE Karras DPMSolverMultistepScheduler init with use_karras_sigmas=True and algorithm_type="sde-dpmsolver++"
|
|
|
|
DPM++ SDE DPMSolverSinglestepScheduler
|
|
DPM++ SDE Karras DPMSolverSinglestepScheduler init with use_karras_sigmas=True
|
|
DPM2 KDPM2DiscreteScheduler
|
|
DPM2 Karras KDPM2DiscreteScheduler init with use_karras_sigmas=True
|
|
DPM2 a KDPM2AncestralDiscreteScheduler
|
|
DPM2 a Karras KDPM2AncestralDiscreteScheduler init with use_karras_sigmas=True
|
|
'''
|
|
# scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
|
|
scheduler = DPMSolverMultistepScheduler()
|
|
elif args.sample_method == 'DPMSolverSinglestep':
|
|
scheduler = DPMSolverSinglestepScheduler()
|
|
elif args.sample_method == 'PNDM':
|
|
scheduler = PNDMScheduler()
|
|
elif args.sample_method == 'HeunDiscrete': ########
|
|
scheduler = HeunDiscreteScheduler()
|
|
elif args.sample_method == 'EulerAncestralDiscrete':
|
|
scheduler = EulerAncestralDiscreteScheduler()
|
|
elif args.sample_method == 'DEISMultistep':
|
|
scheduler = DEISMultistepScheduler()
|
|
elif args.sample_method == 'KDPM2AncestralDiscrete': #########
|
|
scheduler = KDPM2AncestralDiscreteScheduler()
|
|
elif args.sample_method == 'EulerDiscreteSVD':
|
|
scheduler = EulerDiscreteScheduler.from_pretrained("stabilityai/stable-video-diffusion-img2vid",
|
|
subfolder="scheduler", cache_dir=args.cache_dir)
|
|
pipeline = OpenSoraPipeline(vae=vae,
|
|
text_encoder=text_encoder,
|
|
tokenizer=tokenizer,
|
|
scheduler=scheduler,
|
|
transformer=transformer_model)
|
|
pipeline.to(device)
|
|
if args.compile:
|
|
# 5% https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
|
|
options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision", \
|
|
"memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
|
|
"triton.fuse_attention_allow_fp16_reduction": false}}'
|
|
# options = '{"mode": "max-autotune", "memory_format": "channels_last", \
|
|
# "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
|
|
from onediffx import compile_pipe
|
|
pipeline = compile_pipe(
|
|
pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
|
|
)
|
|
|
|
# 4%
|
|
# pipeline.transformer = torch.compile(pipeline.transformer)
|
|
|
|
if not os.path.exists(args.save_img_path):
|
|
os.makedirs(args.save_img_path)
|
|
|
|
if not isinstance(args.text_prompt, list):
|
|
args.text_prompt = [args.text_prompt]
|
|
if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
|
|
text_prompt = open(args.text_prompt[0], 'r').readlines()
|
|
text_prompt = [i.strip() for i in text_prompt]
|
|
|
|
positive_prompt = """
|
|
(masterpiece), (best quality), (ultra-detailed),
|
|
{}.
|
|
"""
|
|
|
|
negative_prompt = """
|
|
nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality,
|
|
low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
|
|
"""
|
|
|
|
# positive_prompt = """
|
|
# (masterpiece), (best quality), (ultra-detailed),
|
|
# {}.
|
|
# emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo,
|
|
# sharp focus, high budget, cinemascope, moody, epic, gorgeous
|
|
# """
|
|
|
|
# negative_prompt = """
|
|
# disfigured, poorly drawn face, longbody, lowres, bad anatomy, bad hands, missing fingers, cropped, worst quality, low quality
|
|
# """
|
|
|
|
video_grids = []
|
|
for idx, prompt in enumerate(text_prompt):
|
|
videos = pipeline(positive_prompt.format(prompt),
|
|
negative_prompt=negative_prompt,
|
|
num_frames=args.num_frames,
|
|
height=args.height,
|
|
width=args.width,
|
|
num_inference_steps=args.num_sampling_steps,
|
|
guidance_scale=args.guidance_scale,
|
|
num_images_per_prompt=1,
|
|
mask_feature=True,
|
|
device=args.device,
|
|
max_sequence_length=args.max_sequence_length,
|
|
).images
|
|
try:
|
|
if args.num_frames == 1:
|
|
ext = 'jpg'
|
|
videos = videos[:, 0].permute(0, 3, 1, 2) # b t h w c -> b c h w
|
|
save_image(videos / 255.0, os.path.join(args.save_img_path, f'{idx}.{ext}'), nrow=1, normalize=True, value_range=(0, 1)) # t c h w
|
|
|
|
else:
|
|
ext = 'mp4'
|
|
imageio.mimwrite(
|
|
os.path.join(args.save_img_path, f'{idx}.{ext}'), videos[0], fps=args.fps, quality=6) # highest quality is 10, lowest is 0
|
|
except:
|
|
print('Error when saving {}'.format(prompt))
|
|
video_grids.append(videos)
|
|
video_grids = torch.cat(video_grids, dim=0)
|
|
|
|
|
|
# torchvision.io.write_video(args.save_img_path + '_%04d' % args.run_time + '-.mp4', video_grids, fps=6)
|
|
if args.num_frames == 1:
|
|
save_image(video_grids / 255.0, os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
|
|
nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
|
|
else:
|
|
video_grids = save_video_grid(video_grids)
|
|
imageio.mimwrite(os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'), video_grids, fps=args.fps, quality=6)
|
|
|
|
print('save path {}'.format(args.save_img_path))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
|
|
parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
|
|
parser.add_argument("--num_frames", type=int, default=1)
|
|
parser.add_argument("--height", type=int, default=512)
|
|
parser.add_argument("--width", type=int, default=512)
|
|
parser.add_argument("--device", type=str, default='cuda:0')
|
|
parser.add_argument("--cache_dir", type=str, default='./cache_dir')
|
|
parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
|
|
parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
|
|
parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
|
|
parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
|
|
parser.add_argument("--guidance_scale", type=float, default=2.5)
|
|
parser.add_argument("--sample_method", type=str, default="PNDM")
|
|
parser.add_argument("--max_sequence_length", type=int, default=300)
|
|
parser.add_argument("--num_sampling_steps", type=int, default=50)
|
|
parser.add_argument("--fps", type=int, default=24)
|
|
parser.add_argument("--run_time", type=int, default=0)
|
|
parser.add_argument("--text_prompt", nargs='+')
|
|
parser.add_argument('--tile_overlap_factor', type=float, default=0.125)
|
|
parser.add_argument('--enable_tiling', action='store_true')
|
|
parser.add_argument('--compile', action='store_true')
|
|
parser.add_argument('--model_type', type=str, default="udit", choices=['dit', 'udit', 'latte'])
|
|
parser.add_argument('--save_memory', action='store_true')
|
|
args = parser.parse_args()
|
|
|
|
main(args) |