227 lines
10 KiB
Python
227 lines
10 KiB
Python
import argparse
|
|
import sys
|
|
import datetime
|
|
import os
|
|
from omegaconf import OmegaConf
|
|
|
|
import torch
|
|
import torchvision.transforms as transforms
|
|
from torchvision.transforms import functional as F
|
|
|
|
import diffusers
|
|
from diffusers import DDIMScheduler
|
|
from diffusers.utils.import_utils import is_xformers_available
|
|
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
|
|
from lvdm.models.unet import UNetModel
|
|
from lvdm.models.autoencoder import AutoencoderKL, AutoencoderKL_Dualref
|
|
from lvdm.models.condition import FrozenOpenCLIPEmbedder, FrozenOpenCLIPImageEmbedderV2, Resampler
|
|
from lvdm.models.layer_controlnet import LayerControlNet
|
|
from lvdm.pipelines.pipeline_animation import AnimationPipeline
|
|
from lvdm.utils import generate_gaussian_heatmap, save_videos_grid, save_videos_with_traj
|
|
|
|
from einops import rearrange
|
|
import decord
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
# import debugpy
|
|
# debugpy.listen(5678)
|
|
# print("Waiting for debugger attach")
|
|
# debugpy.wait_for_client()
|
|
|
|
@torch.no_grad()
|
|
def main(args):
|
|
if args.savedir is None:
|
|
time_str = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
|
|
savedir = f"samples/{Path(args.config).stem}-{time_str}"
|
|
else:
|
|
savedir = args.savedir
|
|
os.makedirs(savedir, exist_ok=True)
|
|
|
|
config = OmegaConf.load(args.config)
|
|
weight_dtype = torch.bfloat16 if config["mixed_precision"] == "bf16" else torch.float32
|
|
mode = config.get("mode", "interpolate")
|
|
# create validation pipeline
|
|
scheduler = DDIMScheduler.from_pretrained(config.pretrained_model_path, subfolder="scheduler")
|
|
text_encoder = FrozenOpenCLIPEmbedder().eval()
|
|
image_encoder = FrozenOpenCLIPImageEmbedderV2().eval()
|
|
image_projector = Resampler.from_pretrained(config.pretrained_model_path, subfolder="image_projector").eval()
|
|
vae, vae_dualref = None, None
|
|
if mode == "interpolate":
|
|
vae_dualref = AutoencoderKL_Dualref.from_pretrained(config.pretrained_model_path, subfolder="vae_dualref").eval()
|
|
else:
|
|
vae = AutoencoderKL.from_pretrained(config.pretrained_model_path, subfolder="vae").eval()
|
|
unet = UNetModel.from_pretrained(config.pretrained_model_path, subfolder="unet").eval()
|
|
layer_controlnet = LayerControlNet.from_pretrained(config.pretrained_model_path, subfolder="layer_controlnet").eval()
|
|
|
|
pipeline = AnimationPipeline(
|
|
vae=vae, vae_dualref=vae_dualref, text_encoder=text_encoder, image_encoder=image_encoder, image_projector=image_projector,
|
|
unet=unet, layer_controlnet=layer_controlnet,
|
|
scheduler=scheduler,
|
|
).to(device=args.device, dtype=weight_dtype)
|
|
if mode == "interpolate":
|
|
pipeline.vae_dualref.decoder.to(dtype=torch.float32)
|
|
if config.enable_xformers_memory_efficient_attention:
|
|
if is_xformers_available():
|
|
pipeline.enable_xformers_memory_efficient_attention()
|
|
else:
|
|
raise ValueError("xformers is not available. Make sure it is installed correctly")
|
|
|
|
if config.seed is None:
|
|
generator = None
|
|
else:
|
|
np.random.seed(config.seed)
|
|
torch.manual_seed(config.seed)
|
|
generator = torch.Generator(args.device).manual_seed(config.seed)
|
|
|
|
config.W = config.get("W", args.W)
|
|
config.H = config.get("H", args.H)
|
|
config.L = config.get("L", args.L)
|
|
|
|
image_transforms = transforms.Compose([
|
|
transforms.Resize(min(config.H, config.W)),
|
|
transforms.CenterCrop((config.H, config.W)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
|
|
])
|
|
mask_transforms = transforms.Compose([
|
|
transforms.Resize(min(config.H, config.W)),
|
|
transforms.CenterCrop((config.H, config.W)),
|
|
])
|
|
|
|
demo_dir = config.get("demo_dir")
|
|
first_frame = os.path.join(demo_dir, "first_frame.jpg")
|
|
if mode == "interpolate":
|
|
last_frame = os.path.join(demo_dir, "last_frame.jpg")
|
|
else:
|
|
last_frame = None
|
|
sketch_path = os.path.join(demo_dir, "sketch.mp4")
|
|
trajectory_path = os.path.join(demo_dir, "trajectory.npz")
|
|
|
|
if last_frame is None:
|
|
# Image to Video
|
|
image = image_transforms(Image.open(first_frame).convert("RGB"))
|
|
frame_tensor = image[None].to(args.device) # [F, C, H, W]
|
|
else:
|
|
# Interpolate
|
|
image1 = image_transforms(Image.open(first_frame).convert("RGB"))
|
|
image2 = image_transforms(Image.open(last_frame).convert("RGB"))
|
|
frame_tensor1 = image1[None]
|
|
frame_tensor2 = image2[None]
|
|
frame_tensor = torch.cat([frame_tensor1, frame_tensor2], dim=0).to(args.device)
|
|
frame_tensor = frame_tensor[None]
|
|
|
|
if mode == "interpolate":
|
|
layer_masks = torch.zeros((1, config.layer_capacity, 2, 1, config.H, config.W), dtype=torch.bool)
|
|
else:
|
|
layer_masks = torch.zeros((1, config.layer_capacity, 1, 1, config.H, config.W), dtype=torch.bool)
|
|
for layer_idx in range(config.layer_capacity):
|
|
mask_path = os.path.join(demo_dir, f"layer_{layer_idx}.jpg")
|
|
if os.path.exists(mask_path):
|
|
mask = mask_transforms(Image.open(mask_path).convert("L"))
|
|
mask = F.to_tensor(mask) > 0.5
|
|
layer_masks[0, layer_idx, 0] = mask
|
|
last_mask_path = os.path.join(demo_dir, f"layer_{layer_idx}_last.jpg")
|
|
if os.path.exists(last_mask_path) and mode == "interpolate":
|
|
mask = mask_transforms(Image.open(last_mask_path).convert("L"))
|
|
mask = F.to_tensor(mask) > 0.5
|
|
layer_masks[0, layer_idx, 1] = mask
|
|
layer_masks = layer_masks.to(args.device)
|
|
layer_regions = layer_masks * frame_tensor[:, None]
|
|
layer_validity = torch.tensor([config.layer_validity], dtype=torch.bool, device=args.device)
|
|
motion_scores = torch.tensor([config.motion_scores], dtype=weight_dtype, device=args.device)
|
|
layer_static = torch.tensor([config.layer_static], dtype=torch.bool, device=args.device)
|
|
|
|
sketch = torch.ones((1, config.layer_capacity, config.L, 3, config.H, config.W), dtype=weight_dtype)
|
|
if os.path.exists(sketch_path):
|
|
video_reader = decord.VideoReader(sketch_path)
|
|
assert len(video_reader) == config.L, f"Input the length of sketch sequence should match the video length."
|
|
video_frames = video_reader.get_batch(range(config.L)).asnumpy()
|
|
sketch_values = [image_transforms(Image.fromarray(frame)) for frame in video_frames]
|
|
sketch_values = torch.stack(sketch_values, dim=0)
|
|
sketch[0, config.sketch_layer_index] = sketch_values
|
|
sketch = sketch.to(args.device)
|
|
|
|
heatmap = torch.zeros((1, config.layer_capacity, config.L, 3, config.H, config.W), dtype=weight_dtype)
|
|
heatmap[:, :, :, 0] -= 1
|
|
if os.path.exists(trajectory_path):
|
|
traj_file = np.load(trajectory_path)
|
|
traj_width = traj_file["width"]
|
|
traj_height = traj_file["height"]
|
|
trajectory = traj_file["trajectory"]
|
|
if traj_width < traj_height:
|
|
scale = min(config.H, config.W) / traj_width
|
|
new_h = int(traj_height * scale)
|
|
new_w = min(config.H, config.W)
|
|
else:
|
|
scale = min(config.H, config.W) / traj_height
|
|
new_w = int(traj_width * scale)
|
|
new_h = min(config.H, config.W)
|
|
trajectory[..., :2] *= scale
|
|
crop_x = int(round((new_w - config.W) / 2.0))
|
|
crop_y = int(round((new_h - config.H) / 2.0))
|
|
trajectory[..., 0] -= crop_x
|
|
trajectory[..., 1] -= crop_y
|
|
traj_layer_index = torch.zeros(trajectory.shape[1], dtype=torch.long) + config.traj_layer_index
|
|
|
|
heatmap = generate_gaussian_heatmap(trajectory, config.W, config.H, traj_layer_index, config.layer_capacity, offset=True)
|
|
heatmap = rearrange(heatmap, "f n c h w -> (f n) c h w")
|
|
graymap, offset = heatmap[:, :1], heatmap[:, 1:]
|
|
graymap = graymap / 255.
|
|
rad = torch.sqrt(offset[:, 0:1]**2 + offset[:, 1:2]**2)
|
|
rad_max = torch.max(rad)
|
|
epsilon = 1e-5
|
|
offset = offset / (rad_max + epsilon)
|
|
graymap = graymap * 2 - 1
|
|
heatmap = torch.cat([graymap, offset], dim=1)
|
|
heatmap = mask_transforms(heatmap) # no need for normalization
|
|
heatmap = rearrange(heatmap, '(f n) c h w -> n f c h w', n=config.layer_capacity)
|
|
heatmap = heatmap[None]
|
|
heatmap = heatmap.to(args.device)
|
|
|
|
sample = pipeline(
|
|
config.prompt,
|
|
config.L,
|
|
config.H,
|
|
config.W,
|
|
frame_tensor,
|
|
layer_masks = layer_masks,
|
|
layer_regions = layer_regions,
|
|
layer_static = layer_static,
|
|
motion_scores = motion_scores,
|
|
sketch = sketch,
|
|
trajectory = heatmap,
|
|
layer_validity = layer_validity,
|
|
num_inference_steps = config.num_inference_steps,
|
|
guidance_scale = config.guidance_scale,
|
|
guidance_rescale = config.guidance_rescale,
|
|
negative_prompt = config.n_prompt,
|
|
num_videos_per_prompt = config.num_videos_per_prompt,
|
|
eta = config.eta,
|
|
generator = generator,
|
|
fps = config.fps,
|
|
mode = mode,
|
|
weight_dtype = weight_dtype,
|
|
output_type = "tensor",
|
|
).videos
|
|
|
|
for idx, video in enumerate(sample):
|
|
save_videos_grid(video[None], os.path.join(savedir, f"video_{idx}.mp4"), fps=8)
|
|
print(f"Saved {os.path.join(savedir, f'video_{idx}.mp4')}")
|
|
if os.path.exists(trajectory_path):
|
|
save_videos_with_traj(video, torch.from_numpy(trajectory), os.path.join(savedir, f"video_{idx}_with_traj.mp4"), fps=8, line_width=7, circle_radius=10)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=str, required=True)
|
|
parser.add_argument("--savedir", type=str, default=None)
|
|
|
|
parser.add_argument("--L", type=int, default=16 )
|
|
parser.add_argument("--W", type=int, default=512)
|
|
parser.add_argument("--H", type=int, default=320)
|
|
parser.add_argument("--device", type=str, default="cuda:0")
|
|
args = parser.parse_args()
|
|
main(args)
|