import argparse import os import warnings from pathlib import Path from uuid import uuid4 from utils.lora import inject_inferable_lora import torch from diffusers import DPMSolverMultistepScheduler, TextToVideoSDPipeline from models.unet_3d_condition import UNet3DConditionModel from einops import rearrange from torch.nn.functional import interpolate import imageio import decord from train import handle_memory_attention, load_primary_models from utils.lama import inpaint_watermark def initialize_pipeline(model, device="cuda", xformers=False, sdp=False): with warnings.catch_warnings(): warnings.simplefilter("ignore") scheduler, tokenizer, text_encoder, vae, _unet = load_primary_models(model) del _unet #This is a no op unet = UNet3DConditionModel.from_pretrained(model, subfolder='unet') # unet.disable_gradient_checkpointing() pipeline = TextToVideoSDPipeline.from_pretrained( pretrained_model_name_or_path=model, scheduler=scheduler, tokenizer=tokenizer, text_encoder=text_encoder.to(device=device, dtype=torch.half), vae=vae.to(device=device, dtype=torch.half), unet=unet.to(device=device, dtype=torch.half), ) pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) unet._set_gradient_checkpointing(value=False) handle_memory_attention(xformers, sdp, unet) vae.enable_slicing() return pipeline def vid2vid( pipeline, init_video, init_weight, prompt, negative_prompt, height, width, num_inference_steps, generator, guidance_scale ): num_frames = init_video.shape[2] init_video = rearrange(init_video, "b c f h w -> (b f) c h w") pipeline.generator=generator latents = pipeline.vae.encode(init_video).latent_dist.sample() latents = rearrange(latents, "(b f) c h w -> b c f h w", f=num_frames) latents = pipeline.scheduler.add_noise( original_samples=latents * 0.18215, noise=torch.randn_like(latents), timesteps=(torch.ones(latents.shape[0]) * pipeline.scheduler.num_train_timesteps * (1 - init_weight)).long(), ) if latents.shape[0] != len(prompt): latents = latents.repeat(len(prompt), 1, 1, 1, 1) do_classifier_free_guidance = guidance_scale > 1.0 prompt_embeds = pipeline._encode_prompt( prompt=prompt, negative_prompt=negative_prompt, device=latents.device, num_images_per_prompt=1, do_classifier_free_guidance=do_classifier_free_guidance, ) pipeline.scheduler.set_timesteps(num_inference_steps, device=latents.device) timesteps = pipeline.scheduler.timesteps timesteps = timesteps[round(init_weight * len(timesteps)) :] with pipeline.progress_bar(total=len(timesteps)) as progress_bar: for t in timesteps: # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t) # predict the noise residual noise_pred = pipeline.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # reshape latents bsz, channel, frames, width, height = latents.shape latents = latents.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height) noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height) # compute the previous noisy sample x_t -> x_t-1 latents = pipeline.scheduler.step(noise_pred, t, latents).prev_sample # reshape latents back latents = latents[None, :].reshape(bsz, frames, channel, width, height).permute(0, 2, 1, 3, 4) progress_bar.update() video_tensor = pipeline.decode_latents(latents) return video_tensor @torch.inference_mode() def inference( model, prompt, negative_prompt=None, batch_size=1, num_frames=16, width=256, height=256, num_steps=50, guidance_scale=9, init_video=None, init_weight=0.5, device="cuda", xformers=False, sdp=False, lora_path='', lora_rank=64, seed=0, ): with torch.autocast(device, dtype=torch.half): pipeline = initialize_pipeline(model, device, xformers, sdp) inject_inferable_lora(pipeline, lora_path, r=lora_rank) prompt = [prompt] * batch_size negative_prompt = ([negative_prompt] * batch_size) if negative_prompt is not None else None if init_video is not None: g_cuda = torch.Generator(device='cuda') g_cuda.manual_seed(seed) g_cpu = torch.Generator() g_cpu.manual_seed(seed) videos = vid2vid( pipeline=pipeline, init_video=init_video.to(device=device, dtype=torch.half), init_weight=init_weight, prompt=prompt, negative_prompt=negative_prompt, height=height, width=width, num_inference_steps=num_steps, generator=g_cuda, guidance_scale=guidance_scale, ) else: g_cuda = torch.Generator(device='cuda') g_cuda.manual_seed(seed) g_cpu = torch.Generator() g_cpu.manual_seed(seed) videos = pipeline( prompt=prompt, negative_prompt=negative_prompt, num_frames=num_frames, height=height, width=width, num_inference_steps=num_steps, generator=g_cuda, guidance_scale=guidance_scale, output_type="pt", ).frames return videos def export_to_video(video_frames, output_video_path, fps): writer = imageio.get_writer(output_video_path, format="FFMPEG", fps=fps) for frame in video_frames: writer.append_data(frame) writer.close() def run(**args): decord.bridge.set_bridge("torch") output_dir = args.pop("output_dir") fps = args.pop("fps") remove_watermark = args.pop("remove_watermark") init_video = args.get("init_video", None) if init_video is not None: vr = decord.VideoReader(init_video) init = rearrange(vr[:], "f h w c -> c f h w").div(127.5).sub(1).unsqueeze(0) init = interpolate(init, size=(args['num_frames'], args['height'], args['width']), mode="trilinear") args["init_video"] = init videos = inference(**args) os.makedirs(output_dir, exist_ok=True) for idx, video in enumerate(videos): if remove_watermark: video = rearrange(video, "c f h w -> f c h w").add(1).div(2) video = inpaint_watermark(video) video = rearrange(video, "f c h w -> f h w c").clamp(0, 1).mul(255) else: video = rearrange(video, "c f h w -> f h w c").clamp(-1, 1).add(1).mul(127.5) video = video.byte().cpu().numpy() filename = os.path.join(output_dir, f"output-{idx}.mp4") export_to_video(video, filename, fps) yield filename if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("-m", "--model", type=str, required=True) parser.add_argument("-p", "--prompt", type=str, required=True) parser.add_argument("-n", "--negative_prompt", type=str, default=None) parser.add_argument("-o", "--output_dir", type=str, default="./output") parser.add_argument("-B", "--batch_size", type=int, default=1) parser.add_argument("-T", "--num_frames", type=int, default=16) parser.add_argument("-W", "--width", type=int, default=256) parser.add_argument("-H", "--height", type=int, default=256) parser.add_argument("-s", "--num_steps", type=int, default=50) parser.add_argument("-g", "--guidance-scale", type=float, default=9) parser.add_argument("-i", "--init-video", type=str, default=None) parser.add_argument("-iw", "--init-weight", type=float, default=0.5) parser.add_argument("-f", "--fps", type=int, default=8) parser.add_argument("-d", "--device", type=str, default="cuda") parser.add_argument("-x", "--xformers", action="store_true") parser.add_argument("-S", "--sdp", action="store_true") parser.add_argument("-lP", "--lora_path", type=str, default="") parser.add_argument("-lR", "--lora_rank", type=int, default=64) parser.add_argument("-rw", "--remove-watermark", action="store_true") parser.add_argument("-seed", "--seed", type=int, default =0) args = vars(parser.parse_args()) for filename in run(**args): print(filename)