diff --git a/.gitignore b/.gitignore index 8623d1a9f..490bba564 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,5 @@ /outputs /build /src -/.vscode \ No newline at end of file +/.vscode +**/__pycache__/ diff --git a/README.md b/README.md index 8454fc9f4..0a9372137 100755 --- a/README.md +++ b/README.md @@ -5,6 +5,46 @@ ## News +**April 4, 2025** +- We are releasing **[Stable Video 4D 2.0 (SV4D 2.0)](https://huggingface.co/stabilityai/sv4d2.0)**, an enhanced video-to-4D diffusion model for high-fidelity novel-view video synthesis and 4D asset generation. For research purposes: + - **SV4D 2.0** was trained to generate 48 frames (12 video frames x 4 camera views) at 576x576 resolution, given a 12-frame input video of the same size, ideally consisting of white-background images of a moving object. + - Compared to our previous 4D model [SV4D](https://huggingface.co/stabilityai/sv4d), **SV4D 2.0** can generate videos with higher fidelity, sharper details during motion, and better spatio-temporal consistency. It also generalizes much better to real-world videos. Moreover, it does not rely on refernce multi-view of the first frame generated by SV3D, making it more robust to self-occlusions. + - To generate longer novel-view videos, we autoregressively generate 12 frames at a time and use the previous generation as conditioning views for the remaining frames. + - Please check our [project page](https://sv4d20.github.io), [arxiv paper](https://arxiv.org/pdf/2503.16396) and [video summary](https://www.youtube.com/watch?v=dtqj-s50ynU) for more details. + +**QUICKSTART** : +- `python scripts/sampling/simple_video_sample_4d2.py --input_path assets/sv4d_videos/camel.gif --output_folder outputs` (after downloading [sv4d2.safetensors](https://huggingface.co/stabilityai/sv4d2.0) from HuggingFace into `checkpoints/`) + +To run **SV4D 2.0** on a single input video of 21 frames: +- Download SV4D 2.0 model (`sv4d2.safetensors`) from [here](https://huggingface.co/stabilityai/sv4d2.0) to `checkpoints/`: `huggingface-cli download stabilityai/sv4d2.0 sv4d2.safetensors --local-dir checkpoints` +- Run inference: `python scripts/sampling/simple_video_sample_4d2.py --input_path ` + - `input_path` : The input video `` can be + - a single video file in `gif` or `mp4` format, such as `assets/sv4d_videos/camel.gif`, or + - a folder containing images of video frames in `.jpg`, `.jpeg`, or `.png` format, or + - a file name pattern matching images of video frames. + - `num_steps` : default is 50, can decrease to it to shorten sampling time. + - `elevations_deg` : specified elevations (reletive to input view), default is 0.0 (same as input view). + - **Background removal** : For input videos with plain background, (optionally) use [rembg](https://github.com/danielgatis/rembg) to remove background and crop video frames by setting `--remove_bg=True`. To obtain higher quality outputs on real-world input videos with noisy background, try segmenting the foreground object using [Clipdrop](https://clipdrop.co/) or [SAM2](https://github.com/facebookresearch/segment-anything-2) before running SV4D. + - **Low VRAM environment** : To run on GPUs with low VRAM, try setting `--encoding_t=1` (of frames encoded at a time) and `--decoding_t=1` (of frames decoded at a time) or lower video resolution like `--img_size=512`. + +Notes: +- We also train a 8-view model that generates 5 frames x 8 views at a time (same as SV4D). + - Download the model from huggingface: `huggingface-cli download stabilityai/sv4d2.0 sv4d2_8views.safetensors --local-dir checkpoints` + - Run inference: `python scripts/sampling/simple_video_sample_4d2.py --model_path checkpoints/sv4d2_8views.safetensors --input_path assets/sv4d_videos/chest.gif --output_folder outputs` + - The 5x8 model takes 5 frames of input at a time. But the inference scripts for both model take 21-frame video as input by default (same as SV3D and SV4D), we run the model autoregressively until we generate 21 frames. +- Install dependencies before running: +``` +python3.10 -m venv .generativemodels +source .generativemodels/bin/activate +pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # check CUDA version +pip3 install -r requirements/pt2.txt +pip3 install . +pip3 install -e git+https://github.com/Stability-AI/datapipelines.git@main#egg=sdata +``` + + ![tile](assets/sv4d2.gif) + + **July 24, 2024** - We are releasing **[Stable Video 4D (SV4D)](https://huggingface.co/stabilityai/sv4d)**, a video-to-4D diffusion model for novel-view video synthesis. For research purposes: - **SV4D** was trained to generate 40 frames (5 video frames x 8 camera views) at 576x576 resolution, given 5 context frames (the input video), and 8 reference views (synthesised from the first frame of the input video, using a multi-view diffusion model like SV3D) of the same size, ideally white-background images with one object. @@ -164,6 +204,7 @@ This is assuming you have navigated to the `generative-models` root after clonin # install required packages from pypi python3 -m venv .pt2 source .pt2/bin/activate +pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip3 install -r requirements/pt2.txt ``` diff --git a/assets/sv4d2.gif b/assets/sv4d2.gif new file mode 100644 index 000000000..04125d3fc Binary files /dev/null and b/assets/sv4d2.gif differ diff --git a/assets/sv4d_videos/bear.gif b/assets/sv4d_videos/bear.gif new file mode 100644 index 000000000..679e42e2c Binary files /dev/null and b/assets/sv4d_videos/bear.gif differ diff --git a/assets/sv4d_videos/bee.gif b/assets/sv4d_videos/bee.gif new file mode 100644 index 000000000..b3a916547 Binary files /dev/null and b/assets/sv4d_videos/bee.gif differ diff --git a/assets/sv4d_videos/bmx-bumps.gif b/assets/sv4d_videos/bmx-bumps.gif new file mode 100644 index 000000000..64730f397 Binary files /dev/null and b/assets/sv4d_videos/bmx-bumps.gif differ diff --git a/assets/sv4d_videos/bunnyman.mp4 b/assets/sv4d_videos/bunnyman.mp4 deleted file mode 100644 index 40a5b27d0..000000000 Binary files a/assets/sv4d_videos/bunnyman.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/camel.gif b/assets/sv4d_videos/camel.gif new file mode 100644 index 000000000..5e8adfa8a Binary files /dev/null and b/assets/sv4d_videos/camel.gif differ diff --git a/assets/sv4d_videos/chameleon.gif b/assets/sv4d_videos/chameleon.gif new file mode 100644 index 000000000..b517450fc Binary files /dev/null and b/assets/sv4d_videos/chameleon.gif differ diff --git a/assets/sv4d_videos/chest.gif b/assets/sv4d_videos/chest.gif new file mode 100644 index 000000000..9f5d8e439 Binary files /dev/null and b/assets/sv4d_videos/chest.gif differ diff --git a/assets/sv4d_videos/cows.gif b/assets/sv4d_videos/cows.gif new file mode 100644 index 000000000..a857c744c Binary files /dev/null and b/assets/sv4d_videos/cows.gif differ diff --git a/assets/sv4d_videos/dance-twirl.gif b/assets/sv4d_videos/dance-twirl.gif new file mode 100644 index 000000000..8a10bce8c Binary files /dev/null and b/assets/sv4d_videos/dance-twirl.gif differ diff --git a/assets/sv4d_videos/dolphin.mp4 b/assets/sv4d_videos/dolphin.mp4 deleted file mode 100644 index e80f4601c..000000000 Binary files a/assets/sv4d_videos/dolphin.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/flag.gif b/assets/sv4d_videos/flag.gif new file mode 100644 index 000000000..a1f3552e7 Binary files /dev/null and b/assets/sv4d_videos/flag.gif differ diff --git a/assets/sv4d_videos/gear.gif b/assets/sv4d_videos/gear.gif new file mode 100644 index 000000000..75d750816 Binary files /dev/null and b/assets/sv4d_videos/gear.gif differ diff --git a/assets/sv4d_videos/green_robot.mp4 b/assets/sv4d_videos/green_robot.mp4 deleted file mode 100644 index 428049c62..000000000 Binary files a/assets/sv4d_videos/green_robot.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/guppie_v0.mp4 b/assets/sv4d_videos/guppie_v0.mp4 deleted file mode 100644 index 76764d5c0..000000000 Binary files a/assets/sv4d_videos/guppie_v0.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/hike.gif b/assets/sv4d_videos/hike.gif new file mode 100644 index 000000000..7632b2bc6 Binary files /dev/null and b/assets/sv4d_videos/hike.gif differ diff --git a/assets/sv4d_videos/hiphop_parrot.mp4 b/assets/sv4d_videos/hiphop_parrot.mp4 deleted file mode 100644 index f86514d42..000000000 Binary files a/assets/sv4d_videos/hiphop_parrot.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/horsejump-low.gif b/assets/sv4d_videos/horsejump-low.gif new file mode 100644 index 000000000..c913135e3 Binary files /dev/null and b/assets/sv4d_videos/horsejump-low.gif differ diff --git a/assets/sv4d_videos/human5.mp4 b/assets/sv4d_videos/human5.mp4 deleted file mode 100644 index 1962b5e69..000000000 Binary files a/assets/sv4d_videos/human5.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/human7.mp4 b/assets/sv4d_videos/human7.mp4 deleted file mode 100644 index 3d8bb72bf..000000000 Binary files a/assets/sv4d_videos/human7.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/lucia_v000.mp4 b/assets/sv4d_videos/lucia_v000.mp4 deleted file mode 100644 index 83f4d9e96..000000000 Binary files a/assets/sv4d_videos/lucia_v000.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/monkey.mp4 b/assets/sv4d_videos/monkey.mp4 deleted file mode 100644 index 5434a246e..000000000 Binary files a/assets/sv4d_videos/monkey.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/pistol_v0.mp4 b/assets/sv4d_videos/pistol_v0.mp4 deleted file mode 100644 index 83c0e9821..000000000 Binary files a/assets/sv4d_videos/pistol_v0.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/robot.gif b/assets/sv4d_videos/robot.gif new file mode 100644 index 000000000..c3b6f5f70 Binary files /dev/null and b/assets/sv4d_videos/robot.gif differ diff --git a/assets/sv4d_videos/snowboard.gif b/assets/sv4d_videos/snowboard.gif new file mode 100644 index 000000000..29209bf0d Binary files /dev/null and b/assets/sv4d_videos/snowboard.gif differ diff --git a/assets/sv4d_videos/snowboard_v000.mp4 b/assets/sv4d_videos/snowboard_v000.mp4 deleted file mode 100644 index 5c1b67a42..000000000 Binary files a/assets/sv4d_videos/snowboard_v000.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/stroller_v000.mp4 b/assets/sv4d_videos/stroller_v000.mp4 deleted file mode 100644 index e293bd5ce..000000000 Binary files a/assets/sv4d_videos/stroller_v000.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/test_video2.mp4 b/assets/sv4d_videos/test_video2.mp4 deleted file mode 100644 index 424f14dca..000000000 Binary files a/assets/sv4d_videos/test_video2.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/train_v0.mp4 b/assets/sv4d_videos/train_v0.mp4 deleted file mode 100644 index cb5f76f32..000000000 Binary files a/assets/sv4d_videos/train_v0.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/wave_hello.mp4 b/assets/sv4d_videos/wave_hello.mp4 deleted file mode 100644 index 4c7693f57..000000000 Binary files a/assets/sv4d_videos/wave_hello.mp4 and /dev/null differ diff --git a/assets/sv4d_videos/windmill.gif b/assets/sv4d_videos/windmill.gif new file mode 100644 index 000000000..64d394eeb Binary files /dev/null and b/assets/sv4d_videos/windmill.gif differ diff --git a/requirements/pt2.txt b/requirements/pt2.txt index 824473abd..6a73eebee 100644 --- a/requirements/pt2.txt +++ b/requirements/pt2.txt @@ -5,13 +5,16 @@ einops>=0.6.1 fairscale>=0.4.13 fire>=0.5.0 fsspec>=2023.6.0 +imageio[ffmpeg] +imageio[pyav] invisible-watermark>=0.2.0 kornia==0.6.9 matplotlib>=3.7.2 natsort>=8.4.0 ninja>=1.11.1 -numpy>=1.24.4 +numpy==2.1 omegaconf>=2.3.0 +onnxruntime open-clip-torch>=2.20.0 opencv-python==4.6.0.66 pandas>=2.0.3 diff --git a/scripts/demo/sv4d_helpers.py b/scripts/demo/sv4d_helpers.py index 7296c58c2..d64266c93 100755 --- a/scripts/demo/sv4d_helpers.py +++ b/scripts/demo/sv4d_helpers.py @@ -13,9 +13,6 @@ from omegaconf import ListConfig, OmegaConf from PIL import Image, ImageSequence from rembg import remove -from torch import autocast -from torchvision.transforms import ToTensor - from scripts.util.detection.nsfw_and_watermark_dectection import DeepFloydDataFiltering from sgm.modules.autoencoding.temporal_ae import VideoDecoder from sgm.modules.diffusionmodules.guiders import ( @@ -34,6 +31,8 @@ LinearMultistepSampler, ) from sgm.util import default, instantiate_from_config +from torch import autocast +from torchvision.transforms import ToTensor def load_module_gpu(model): @@ -165,7 +164,16 @@ def read_video( return images_v0 -def preprocess_video(input_path, remove_bg=False, n_frames=21, W=576, H=576, output_folder=None, image_frame_ratio = 0.917): +def preprocess_video( + input_path, + remove_bg=False, + n_frames=21, + W=576, + H=576, + output_folder=None, + image_frame_ratio=0.917, + base_count=0, +): print(f"preprocess {input_path}") if output_folder is None: output_folder = os.path.dirname(input_path) @@ -199,7 +207,9 @@ def preprocess_video(input_path, remove_bg=False, n_frames=21, W=576, H=576, out images = [Image.open(img_path) for img_path in all_img_paths] if len(images) != n_frames: - raise ValueError(f"Input video contains {len(images)} frames, fewer than {n_frames} frames.") + raise ValueError( + f"Input video contains {len(images)} frames, fewer than {n_frames} frames." + ) # Remove background for i, image in enumerate(images): @@ -226,18 +236,28 @@ def preprocess_video(input_path, remove_bg=False, n_frames=21, W=576, H=576, out else: # assume the input image has white background ret, mask = cv2.threshold( - (np.array(image).mean(-1) <= white_thresh).astype(np.uint8) * 255, 0, 255, cv2.THRESH_BINARY + (np.array(image).mean(-1) <= white_thresh).astype(np.uint8) * 255, + 0, + 255, + cv2.THRESH_BINARY, ) - + x, y, w, h = cv2.boundingRect(mask) box_coord[0] = min(box_coord[0], x) box_coord[1] = min(box_coord[1], y) box_coord[2] = max(box_coord[2], x + w) box_coord[3] = max(box_coord[3], y + h) - box_square = max(original_center[0] - box_coord[0], original_center[1] - box_coord[1]) + box_square = max( + original_center[0] - box_coord[0], original_center[1] - box_coord[1] + ) box_square = max(box_square, box_coord[2] - original_center[0]) box_square = max(box_square, box_coord[3] - original_center[1]) - x, y, w, h = original_center[0] - box_square, original_center[1] - box_square, 2 * box_square, 2 * box_square + x, y = max(0, original_center[0] - box_square), max( + 0, original_center[1] - box_square + ) + w, h = min(image_arr.shape[0], 2 * box_square), min( + image_arr.shape[1], 2 * box_square + ) box_size = box_square * 2 for image in images: @@ -245,15 +265,15 @@ def preprocess_video(input_path, remove_bg=False, n_frames=21, W=576, H=576, out image = image.convert("RGBA") image_arr = np.array(image) side_len = ( - int(box_size / image_frame_ratio) - if image_frame_ratio is not None - else in_w + int(box_size / image_frame_ratio) if image_frame_ratio is not None else in_w ) padded_image = np.zeros((side_len, side_len, 4), dtype=np.uint8) center = side_len // 2 + box_size_w = min(w, box_size) + box_size_h = min(h, box_size) padded_image[ - center - box_size // 2 : center - box_size // 2 + box_size, - center - box_size // 2 : center - box_size // 2 + box_size, + center - box_size_w // 2 : center - box_size_w // 2 + box_size_w, + center - box_size_h // 2 : center - box_size_h // 2 + box_size_h, ] = image_arr[x : x + w, y : y + h] rgba = Image.fromarray(padded_image).resize((W, H), Image.LANCZOS) @@ -261,14 +281,14 @@ def preprocess_video(input_path, remove_bg=False, n_frames=21, W=576, H=576, out rgba_arr = np.array(rgba) / 255.0 rgb = rgba_arr[..., :3] * rgba_arr[..., -1:] + (1 - rgba_arr[..., -1:]) image = (rgb * 255).astype(np.uint8) - + images_v0.append(image) - - base_count = len(glob(os.path.join(output_folder, "*.mp4"))) // 12 + processed_file = os.path.join(output_folder, f"{base_count:06d}_process_input.mp4") imageio.mimwrite(processed_file, images_v0, fps=10) return processed_file + def sample_sv3d( image, num_frames: Optional[int] = None, # 21 for SV3D @@ -326,6 +346,7 @@ def sample_sv3d( with torch.no_grad(): with torch.autocast(device): + load_module_gpu(model.conditioner) batch, batch_uc = get_batch_sv3d( get_unique_embedder_keys_from_conditioner(model.conditioner), value_dict, @@ -341,6 +362,7 @@ def sample_sv3d( "cond_frames_without_noise", ], ) + unload_module_gpu(model.conditioner) for k in ["crossattn", "concat"]: uc[k] = repeat(uc[k], "b ... -> b t ...", t=num_frames) @@ -361,11 +383,17 @@ def denoiser(input, sigma, c): model.model, input, sigma, c, **additional_model_inputs ) + load_module_gpu(model.model) + load_module_gpu(model.denoiser) samples_z = model.sampler(denoiser, randn, cond=c, uc=uc) - unload_module_gpu(model.model) unload_module_gpu(model.denoiser) + unload_module_gpu(model.model) + + load_module_gpu(model.first_stage_model) model.en_and_decode_n_samples_a_time = decoding_t samples_x = model.decode_first_stage(samples_z) + unload_module_gpu(model.first_stage_model) + samples_x[-1:] = value_dict["cond_frames_without_noise"] samples = torch.clamp(samples_x, min=-1.0, max=1.0) @@ -373,13 +401,17 @@ def denoiser(input, sigma, c): return samples -def decode_latents(model, samples_z, img_matrix, frame_indices, view_indices, timesteps): +def decode_latents( + model, samples_z, img_matrix, frame_indices, view_indices, timesteps +): load_module_gpu(model.first_stage_model) for t in frame_indices: for v in view_indices: - if t != 0 and v != 0: + if True: # t != 0 and v != 0: if isinstance(model.first_stage_model.decoder, VideoDecoder): - samples_x = model.decode_first_stage(samples_z[t, v][None], timesteps=timesteps) + samples_x = model.decode_first_stage( + samples_z[t, v][None], timesteps=timesteps + ) else: samples_x = model.decode_first_stage(samples_z[t, v][None]) samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0) @@ -555,12 +587,15 @@ def get_guider_no_st(options, key): } elif guider == "SpatiotemporalPredictionGuider": max_scale = options.get("cfg", 1.5) + min_scale = options.get("min_cfg", 1.0) guider_config = { "target": "sgm.modules.diffusionmodules.guiders.SpatiotemporalPredictionGuider", "params": { "max_scale": max_scale, + "min_scale": min_scale, "num_frames": options["num_frames"], + "num_views": options["num_views"], **additional_guider_kwargs, }, } @@ -652,7 +687,7 @@ def init_sampling_no_st( options = {} if options is None else options num_rows, num_cols = 1, 1 - steps = options.get("num_steps", 40) + steps = options.get("num_steps", 50) sampler = [ "EulerEDMSampler", "HeunEDMSampler", @@ -688,6 +723,7 @@ def run_img2vid( cond_motion=None, cond_view=None, decoding_t=None, + cond_mv=True, ): options = version_dict["options"] H = version_dict["H"] @@ -714,7 +750,10 @@ def run_img2vid( value_dict["is_image"] = 0 value_dict["is_webvid"] = 0 - value_dict["image_only_indicator"] = 0 + if cond_mv: + value_dict["image_only_indicator"] = 1.0 + else: + value_dict["image_only_indicator"] = 0.0 cond_aug = 0.00 if cond_motion is not None: @@ -722,8 +761,6 @@ def run_img2vid( value_dict["cond_frames"] = ( cond_motion[:, None].repeat(1, cond_view.shape[0], 1, 1, 1).flatten(0, 1) ) - value_dict["cond_motion"] = cond_motion - value_dict["cond_view"] = cond_view else: value_dict["cond_frames_without_noise"] = image value_dict["cond_frames"] = image + cond_aug * torch.randn_like(image) @@ -760,46 +797,112 @@ def run_img2vid( return samples -def prepare_inputs(frame_indices, img_matrix, v0, view_indices, model, version_dict, seed, polars, azims): - load_module_gpu(model.conditioner) +def prepare_inputs_forward_backward( + img_matrix, + view_indices, + frame_indices, + v0, + t0, + t1, + model, + version_dict, + seed, + polars, + azims, +): + # forward sampling + forward_frame_indices = frame_indices.copy() + image = img_matrix[t0][v0] + cond_motion = torch.cat([img_matrix[t][v0] for t in forward_frame_indices], 0) + cond_view = torch.cat([img_matrix[t0][v] for v in view_indices], 0) + forward_inputs = prepare_sampling( + version_dict, + model, + image, + seed, + polars, + azims, + cond_motion, + cond_view, + ) + + # backward sampling + backward_frame_indices = frame_indices[::-1].copy() + image = img_matrix[t1][v0] + cond_motion = torch.cat([img_matrix[t][v0] for t in backward_frame_indices], 0) + cond_view = torch.cat([img_matrix[t1][v] for v in view_indices], 0) + backward_inputs = prepare_sampling( + version_dict, + model, + image, + seed, + polars, + azims, + cond_motion, + cond_view, + ) + return ( + forward_inputs, + forward_frame_indices, + backward_inputs, + backward_frame_indices, + ) + +def prepare_inputs( + frame_indices, + img_matrix, + v0, + view_indices, + model, + version_dict, + seed, + polars, + azims, +): + load_module_gpu(model.conditioner) + # forward sampling forward_frame_indices = frame_indices.copy() t0 = forward_frame_indices[0] image = img_matrix[t0][v0] cond_motion = torch.cat([img_matrix[t][v0] for t in forward_frame_indices], 0) cond_view = torch.cat([img_matrix[t0][v] for v in view_indices], 0) forward_inputs = prepare_sampling( - version_dict, - model, - image, - seed, - polars, - azims, - cond_motion, - cond_view, - ) - + version_dict, + model, + image, + seed, + polars, + azims, + cond_motion, + cond_view, + ) + # backward sampling - backward_frame_indices = frame_indices[ - ::-1 - ].copy() + backward_frame_indices = frame_indices[::-1].copy() t0 = backward_frame_indices[0] image = img_matrix[t0][v0] cond_motion = torch.cat([img_matrix[t][v0] for t in backward_frame_indices], 0) cond_view = torch.cat([img_matrix[t0][v] for v in view_indices], 0) backward_inputs = prepare_sampling( - version_dict, - model, - image, - seed, - polars, - azims, - cond_motion, - cond_view, - ) + version_dict, + model, + image, + seed, + polars, + azims, + cond_motion, + cond_view, + ) unload_module_gpu(model.conditioner) - return forward_inputs, forward_frame_indices, backward_inputs, backward_frame_indices + return ( + forward_inputs, + forward_frame_indices, + backward_inputs, + backward_frame_indices, + ) + def do_sample( model, @@ -854,6 +957,10 @@ def do_sample( lambda y: y[k][: math.prod(num_samples)].to("cuda"), (c, uc) ) + if value_dict["image_only_indicator"] == 0: + c["cond_view"] *= 0 + uc["cond_view"] *= 0 + additional_model_inputs = {} for k in batch2model_input: if k == "image_only_indicator": @@ -869,9 +976,12 @@ def do_sample( SpatiotemporalPredictionGuider, ), ): - additional_model_inputs[k] = torch.zeros( - num_samples[0] * 2, num_samples[1] - ).to("cuda") + additional_model_inputs[k] = ( + torch.zeros(num_samples[0] * 2, num_samples[1]).to( + "cuda" + ) + + value_dict["image_only_indicator"] + ) else: additional_model_inputs[k] = torch.zeros(num_samples).to( "cuda" @@ -886,11 +996,13 @@ def denoiser(input, sigma, c): return model.denoiser( model.model, input, sigma, c, **additional_model_inputs ) + load_module_gpu(model.model) load_module_gpu(model.denoiser) samples_z = sampler(denoiser, randn, cond=c, uc=uc) - unload_module_gpu(model.model) unload_module_gpu(model.denoiser) + unload_module_gpu(model.model) + load_module_gpu(model.first_stage_model) if isinstance(model.first_stage_model.decoder, VideoDecoder): samples_x = model.decode_first_stage( @@ -900,11 +1012,13 @@ def denoiser(input, sigma, c): samples_x = model.decode_first_stage(samples_z) samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0) unload_module_gpu(model.first_stage_model) + if filter is not None: samples = filter(samples) if return_latents: return samples, samples_z + return samples @@ -931,6 +1045,7 @@ def prepare_sampling_( num_samples = [num_samples, T] else: num_samples = [num_samples] + load_module_gpu(model.conditioner) batch, batch_uc = get_batch( get_unique_embedder_keys_from_conditioner(model.conditioner), value_dict, @@ -944,6 +1059,8 @@ def prepare_sampling_( force_uc_zero_embeddings=force_uc_zero_embeddings, force_cond_zero_embeddings=force_cond_zero_embeddings, ) + unload_module_gpu(model.conditioner) + for k in c: if not k == "crossattn": c[k], uc[k] = map( @@ -965,19 +1082,25 @@ def prepare_sampling_( SpatiotemporalPredictionGuider, ), ): - additional_model_inputs[k] = torch.zeros( - num_samples[0] * 2, num_samples[1] - ).to("cuda") + additional_model_inputs[k] = ( + torch.zeros(num_samples[0] * 2, num_samples[1]).to( + "cuda" + ) + + value_dict["image_only_indicator"] + ) else: additional_model_inputs[k] = torch.zeros(num_samples).to( "cuda" ) else: additional_model_inputs[k] = batch[k] + return c, uc, additional_model_inputs -def do_sample_per_step(model, sampler, noisy_latents, c, uc, step, additional_model_inputs): +def do_sample_per_step( + model, sampler, noisy_latents, c, uc, step, additional_model_inputs +): precision_scope = autocast with torch.no_grad(): with precision_scope("cuda"): @@ -1015,6 +1138,8 @@ def denoiser(input, sigma, c): uc, gamma, ) + unload_module_gpu(model.denoiser) + unload_module_gpu(model.model) return samples_z @@ -1053,7 +1178,7 @@ def prepare_sampling( value_dict["is_image"] = 0 value_dict["is_webvid"] = 0 - value_dict["image_only_indicator"] = 0 + value_dict["image_only_indicator"] = 1.0 cond_aug = 0.00 if cond_motion is not None: @@ -1061,8 +1186,6 @@ def prepare_sampling( value_dict["cond_frames"] = ( cond_motion[:, None].repeat(1, cond_view.shape[0], 1, 1, 1).flatten(0, 1) ) - value_dict["cond_motion"] = cond_motion - value_dict["cond_view"] = cond_view else: value_dict["cond_frames_without_noise"] = image value_dict["cond_frames"] = image + cond_aug * torch.randn_like(image) @@ -1073,8 +1196,6 @@ def prepare_sampling( value_dict["cond_motion"] = cond_motion value_dict["cond_view"] = cond_view - # seed_everything(seed) - options["num_frames"] = T sampler, num_rows, num_cols = init_sampling_no_st(options=options) num_samples = num_rows * num_cols @@ -1269,6 +1390,7 @@ def load_model( num_frames: int, num_steps: int, verbose: bool = False, + ckpt_path: str = None, ): config = OmegaConf.load(config) if device == "cuda": @@ -1281,6 +1403,8 @@ def load_model( config.model.params.sampler_config.params.guider_config.params.num_frames = ( num_frames ) + if ckpt_path is not None: + config.model.params.ckpt_path = ckpt_path if device == "cuda": with torch.device(device): model = instantiate_from_config(config.model).to(device).eval() diff --git a/scripts/sampling/configs/sv4d.yaml b/scripts/sampling/configs/sv4d.yaml index 56e6ddb38..3f1f8aae8 100755 --- a/scripts/sampling/configs/sv4d.yaml +++ b/scripts/sampling/configs/sv4d.yaml @@ -23,6 +23,7 @@ model: attention_resolutions: [4, 2, 1] channel_mult: [1, 2, 4, 4] context_dim: 1024 + motion_context_dim: 4 extra_ff_mix_layer: True in_channels: 8 legacy: False diff --git a/scripts/sampling/configs/sv4d2.yaml b/scripts/sampling/configs/sv4d2.yaml new file mode 100755 index 000000000..f55d899c4 --- /dev/null +++ b/scripts/sampling/configs/sv4d2.yaml @@ -0,0 +1,208 @@ +N_TIME: 12 +N_VIEW: 4 +N_FRAMES: 48 + +model: + target: sgm.models.diffusion.DiffusionEngine + params: + scale_factor: 0.18215 + en_and_decode_n_samples_a_time: 8 + disable_first_stage_autocast: True + ckpt_path: checkpoints/sv4d2.safetensors + denoiser_config: + target: sgm.modules.diffusionmodules.denoiser.Denoiser + params: + scaling_config: + target: sgm.modules.diffusionmodules.denoiser_scaling.VScalingWithEDMcNoise + + network_config: + target: sgm.modules.diffusionmodules.video_model.SpatialUNetModelWithTime + params: + adm_in_channels: 1280 + attention_resolutions: [4, 2, 1] + channel_mult: [1, 2, 4, 4] + context_dim: 1024 + motion_context_dim: 4 + extra_ff_mix_layer: True + in_channels: 8 + legacy: False + model_channels: 320 + num_classes: sequential + num_head_channels: 64 + num_res_blocks: 2 + out_channels: 4 + replicate_time_mix_bug: True + spatial_transformer_attn_type: softmax-xformers + time_block_merge_factor: 0.0 + time_block_merge_strategy: learned_with_images + time_kernel_size: [3, 1, 1] + time_mix_legacy: False + transformer_depth: 1 + use_checkpoint: False + use_linear_in_transformer: True + use_spatial_context: True + use_spatial_transformer: True + separate_motion_merge_factor: True + use_motion_attention: True + use_3d_attention: True + use_camera_emb: True + + conditioner_config: + target: sgm.modules.GeneralConditioner + params: + emb_models: + + - input_key: cond_frames_without_noise + target: sgm.modules.encoders.modules.FrozenOpenCLIPImagePredictionEmbedder + is_trainable: False + params: + n_cond_frames: ${N_TIME} + n_copies: 1 + open_clip_embedding_config: + target: sgm.modules.encoders.modules.FrozenOpenCLIPImageEmbedder + params: + freeze: True + + - input_key: cond_frames + target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder + is_trainable: False + params: + is_ae: True + n_cond_frames: ${N_FRAMES} + n_copies: 1 + encoder_config: + target: sgm.models.autoencoder.AutoencoderKLModeOnly + params: + ddconfig: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + embed_dim: 4 + lossconfig: + target: torch.nn.Identity + monitor: val/rec_loss + sigma_cond_config: + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 + sigma_sampler_config: + target: sgm.modules.diffusionmodules.sigma_sampling.ZeroSampler + + - input_key: polar_rad + is_trainable: False + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 512 + + - input_key: azimuth_rad + is_trainable: False + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 512 + + - input_key: cond_view + is_trainable: False + target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder + params: + is_ae: True + n_cond_frames: ${N_VIEW} + n_copies: 1 + encoder_config: + target: sgm.models.autoencoder.AutoencoderKLModeOnly + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + lossconfig: + target: torch.nn.Identity + sigma_sampler_config: + target: sgm.modules.diffusionmodules.sigma_sampling.ZeroSampler + + - input_key: cond_motion + is_trainable: False + target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder + params: + is_ae: True + n_cond_frames: ${N_TIME} + n_copies: 1 + encoder_config: + target: sgm.models.autoencoder.AutoencoderKLModeOnly + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + lossconfig: + target: torch.nn.Identity + sigma_sampler_config: + target: sgm.modules.diffusionmodules.sigma_sampling.ZeroSampler + + first_stage_config: + target: sgm.models.autoencoder.AutoencodingEngine + params: + loss_config: + target: torch.nn.Identity + regularizer_config: + target: sgm.modules.autoencoding.regularizers.DiagonalGaussianRegularizer + encoder_config: + target: torch.nn.Identity + decoder_config: + target: sgm.modules.diffusionmodules.model.Decoder + params: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + + sampler_config: + target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler + params: + num_steps: 50 + discretization_config: + target: sgm.modules.diffusionmodules.discretizer.EDMDiscretization + params: + sigma_max: 500.0 + guider_config: + target: sgm.modules.diffusionmodules.guiders.SpatiotemporalPredictionGuider + params: + max_scale: 1.5 + min_scale: 1.5 + num_frames: ${N_FRAMES} + num_views: ${N_VIEW} + additional_cond_keys: [ cond_view, cond_motion ] diff --git a/scripts/sampling/configs/sv4d2_8views.yaml b/scripts/sampling/configs/sv4d2_8views.yaml new file mode 100755 index 000000000..15fd38aeb --- /dev/null +++ b/scripts/sampling/configs/sv4d2_8views.yaml @@ -0,0 +1,208 @@ +N_TIME: 5 +N_VIEW: 8 +N_FRAMES: 40 + +model: + target: sgm.models.diffusion.DiffusionEngine + params: + scale_factor: 0.18215 + en_and_decode_n_samples_a_time: 8 + disable_first_stage_autocast: True + ckpt_path: checkpoints/sv4d2_8views.safetensors + denoiser_config: + target: sgm.modules.diffusionmodules.denoiser.Denoiser + params: + scaling_config: + target: sgm.modules.diffusionmodules.denoiser_scaling.VScalingWithEDMcNoise + + network_config: + target: sgm.modules.diffusionmodules.video_model.SpatialUNetModelWithTime + params: + adm_in_channels: 1280 + attention_resolutions: [4, 2, 1] + channel_mult: [1, 2, 4, 4] + context_dim: 1024 + motion_context_dim: 4 + extra_ff_mix_layer: True + in_channels: 8 + legacy: False + model_channels: 320 + num_classes: sequential + num_head_channels: 64 + num_res_blocks: 2 + out_channels: 4 + replicate_time_mix_bug: True + spatial_transformer_attn_type: softmax-xformers + time_block_merge_factor: 0.0 + time_block_merge_strategy: learned_with_images + time_kernel_size: [3, 1, 1] + time_mix_legacy: False + transformer_depth: 1 + use_checkpoint: False + use_linear_in_transformer: True + use_spatial_context: True + use_spatial_transformer: True + separate_motion_merge_factor: True + use_motion_attention: True + use_3d_attention: False + use_camera_emb: True + + conditioner_config: + target: sgm.modules.GeneralConditioner + params: + emb_models: + + - input_key: cond_frames_without_noise + target: sgm.modules.encoders.modules.FrozenOpenCLIPImagePredictionEmbedder + is_trainable: False + params: + n_cond_frames: ${N_TIME} + n_copies: 1 + open_clip_embedding_config: + target: sgm.modules.encoders.modules.FrozenOpenCLIPImageEmbedder + params: + freeze: True + + - input_key: cond_frames + target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder + is_trainable: False + params: + is_ae: True + n_cond_frames: ${N_FRAMES} + n_copies: 1 + encoder_config: + target: sgm.models.autoencoder.AutoencoderKLModeOnly + params: + ddconfig: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + embed_dim: 4 + lossconfig: + target: torch.nn.Identity + monitor: val/rec_loss + sigma_cond_config: + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 + sigma_sampler_config: + target: sgm.modules.diffusionmodules.sigma_sampling.ZeroSampler + + - input_key: polar_rad + is_trainable: False + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 512 + + - input_key: azimuth_rad + is_trainable: False + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 512 + + - input_key: cond_view + is_trainable: False + target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder + params: + is_ae: True + n_cond_frames: ${N_VIEW} + n_copies: 1 + encoder_config: + target: sgm.models.autoencoder.AutoencoderKLModeOnly + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + lossconfig: + target: torch.nn.Identity + sigma_sampler_config: + target: sgm.modules.diffusionmodules.sigma_sampling.ZeroSampler + + - input_key: cond_motion + is_trainable: False + target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder + params: + is_ae: True + n_cond_frames: ${N_TIME} + n_copies: 1 + encoder_config: + target: sgm.models.autoencoder.AutoencoderKLModeOnly + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + lossconfig: + target: torch.nn.Identity + sigma_sampler_config: + target: sgm.modules.diffusionmodules.sigma_sampling.ZeroSampler + + first_stage_config: + target: sgm.models.autoencoder.AutoencodingEngine + params: + loss_config: + target: torch.nn.Identity + regularizer_config: + target: sgm.modules.autoencoding.regularizers.DiagonalGaussianRegularizer + encoder_config: + target: torch.nn.Identity + decoder_config: + target: sgm.modules.diffusionmodules.model.Decoder + params: + attn_resolutions: [] + attn_type: vanilla-xformers + ch: 128 + ch_mult: [1, 2, 4, 4] + double_z: True + dropout: 0.0 + in_channels: 3 + num_res_blocks: 2 + out_ch: 3 + resolution: 256 + z_channels: 4 + + sampler_config: + target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler + params: + num_steps: 50 + discretization_config: + target: sgm.modules.diffusionmodules.discretizer.EDMDiscretization + params: + sigma_max: 500.0 + guider_config: + target: sgm.modules.diffusionmodules.guiders.SpatiotemporalPredictionGuider + params: + max_scale: 2.0 + min_scale: 1.5 + num_frames: ${N_FRAMES} + num_views: ${N_VIEW} + additional_cond_keys: [ cond_view, cond_motion ] diff --git a/scripts/sampling/simple_video_sample.py b/scripts/sampling/simple_video_sample.py index 29a8b8581..904db755a 100644 --- a/scripts/sampling/simple_video_sample.py +++ b/scripts/sampling/simple_video_sample.py @@ -163,7 +163,7 @@ def sample( else: with Image.open(input_img_path) as image: if image.mode == "RGBA": - input_image = image.convert("RGB") + image = image.convert("RGB") w, h = image.size if h % 64 != 0 or w % 64 != 0: @@ -172,7 +172,8 @@ def sample( print( f"WARNING: Your image is of size {h}x{w} which is not divisible by 64. We are resizing to {height}x{width}!" ) - + input_image = np.array(image) + image = ToTensor()(input_image) image = image * 2.0 - 1.0 diff --git a/scripts/sampling/simple_video_sample_4d.py b/scripts/sampling/simple_video_sample_4d.py index 704d7a29a..2a9e15f34 100755 --- a/scripts/sampling/simple_video_sample_4d.py +++ b/scripts/sampling/simple_video_sample_4d.py @@ -27,7 +27,7 @@ def sample( - input_path: str = "assets/test_video.mp4", # Can either be image file or folder with image files + input_path: str = "assets/sv4d_videos/test_video1.mp4", # Can either be image file or folder with image files output_folder: Optional[str] = "outputs/sv4d", num_steps: Optional[int] = 20, sv3d_version: str = "sv3d_u", # sv3d_u or sv3d_p @@ -71,7 +71,8 @@ def sample( "f": F, "options": { "discretization": 1, - "cfg": 3.0, + "cfg": 2.0, + "num_views": V, "sigma_min": 0.002, "sigma_max": 700.0, "rho": 7.0, @@ -94,6 +95,7 @@ def sample( # Read input video frames i.e. images at view 0 print(f"Reading {input_path}") + base_count = len(glob(os.path.join(output_folder, "*.mp4"))) // 11 processed_input_path = preprocess_video( input_path, remove_bg=remove_bg, @@ -102,6 +104,7 @@ def sample( H=H, output_folder=output_folder, image_frame_ratio=image_frame_ratio, + base_count=base_count, ) images_v0 = read_video(processed_input_path, n_frames=n_frames, device=device) @@ -145,15 +148,14 @@ def sample( for t in range(n_frames): img_matrix[t][0] = images_v0[t] - base_count = len(glob(os.path.join(output_folder, "*.mp4"))) // 12 save_video( os.path.join(output_folder, f"{base_count:06d}_t000.mp4"), img_matrix[0], ) - save_video( - os.path.join(output_folder, f"{base_count:06d}_v000.mp4"), - [img_matrix[t][0] for t in range(n_frames)], - ) + # save_video( + # os.path.join(output_folder, f"{base_count:06d}_v000.mp4"), + # [img_matrix[t][0] for t in range(n_frames)], + # ) # Load SV4D model model, filter = load_model( diff --git a/scripts/sampling/simple_video_sample_4d2.py b/scripts/sampling/simple_video_sample_4d2.py new file mode 100755 index 000000000..269c9bfa6 --- /dev/null +++ b/scripts/sampling/simple_video_sample_4d2.py @@ -0,0 +1,235 @@ +import os +import sys +from glob import glob +from typing import List, Optional + +from tqdm import tqdm + +sys.path.append(os.path.realpath(os.path.join(os.path.dirname(__file__), "../../"))) +import numpy as np +import torch +from fire import Fire +from scripts.demo.sv4d_helpers import ( + load_model, + preprocess_video, + read_video, + run_img2vid, + save_video, +) +from sgm.modules.encoders.modules import VideoPredictionEmbedderWithEncoder + +sv4d2_configs = { + "sv4d2": { + "T": 12, # number of frames per sample + "V": 4, # number of views per sample + "model_config": "scripts/sampling/configs/sv4d2.yaml", + "version_dict": { + "T": 12 * 4, + "options": { + "discretization": 1, + "cfg": 2.0, + "min_cfg": 2.0, + "num_views": 4, + "sigma_min": 0.002, + "sigma_max": 700.0, + "rho": 7.0, + "guider": 2, + "force_uc_zero_embeddings": [ + "cond_frames", + "cond_frames_without_noise", + "cond_view", + "cond_motion", + ], + "additional_guider_kwargs": { + "additional_cond_keys": ["cond_view", "cond_motion"] + }, + }, + }, + }, + "sv4d2_8views": { + "T": 5, # number of frames per sample + "V": 8, # number of views per sample + "model_config": "scripts/sampling/configs/sv4d2_8views.yaml", + "version_dict": { + "T": 5 * 8, + "options": { + "discretization": 1, + "cfg": 2.5, + "min_cfg": 1.5, + "num_views": 8, + "sigma_min": 0.002, + "sigma_max": 700.0, + "rho": 7.0, + "guider": 5, + "force_uc_zero_embeddings": [ + "cond_frames", + "cond_frames_without_noise", + "cond_view", + "cond_motion", + ], + "additional_guider_kwargs": { + "additional_cond_keys": ["cond_view", "cond_motion"] + }, + }, + }, + }, +} + + +def sample( + input_path: str = "assets/sv4d_videos/camel.gif", # Can either be image file or folder with image files + model_path: Optional[str] = "checkpoints/sv4d2.safetensors", + output_folder: Optional[str] = "outputs", + num_steps: Optional[int] = 50, + img_size: int = 576, # image resolution + n_frames: int = 21, # number of input and output video frames + seed: int = 23, + encoding_t: int = 8, # Number of frames encoded at a time! This eats most VRAM. Reduce if necessary. + decoding_t: int = 4, # Number of frames decoded at a time! This eats most VRAM. Reduce if necessary. + device: str = "cuda", + elevations_deg: Optional[List[float]] = 0.0, + azimuths_deg: Optional[List[float]] = None, + image_frame_ratio: Optional[float] = 0.9, + verbose: Optional[bool] = False, + remove_bg: bool = False, +): + """ + Simple script to generate multiple novel-view videos conditioned on a video `input_path` or multiple frames, one for each + image file in folder `input_path`. If you run out of VRAM, try decreasing `decoding_t` and `encoding_t`. + """ + # Set model config + assert os.path.basename(model_path) in [ + "sv4d2.safetensors", + "sv4d2_8views.safetensors", + ] + sv4d2_model = os.path.splitext(os.path.basename(model_path))[0] + config = sv4d2_configs[sv4d2_model] + print(sv4d2_model, config) + T = config["T"] + V = config["V"] + model_config = config["model_config"] + version_dict = config["version_dict"] + F = 8 # vae factor to downsize image->latent + C = 4 + H, W = img_size, img_size + n_views = V + 1 # number of output video views (1 input view + 8 novel views) + subsampled_views = np.arange(n_views) + version_dict["H"] = H + version_dict["W"] = W + version_dict["C"] = C + version_dict["f"] = F + version_dict["options"]["num_steps"] = num_steps + + torch.manual_seed(seed) + output_folder = os.path.join(output_folder, sv4d2_model) + os.makedirs(output_folder, exist_ok=True) + + # Read input video frames i.e. images at view 0 + print(f"Reading {input_path}") + base_count = len(glob(os.path.join(output_folder, "*.mp4"))) // n_views + processed_input_path = preprocess_video( + input_path, + remove_bg=remove_bg, + n_frames=n_frames, + W=W, + H=H, + output_folder=output_folder, + image_frame_ratio=image_frame_ratio, + base_count=base_count, + ) + images_v0 = read_video(processed_input_path, n_frames=n_frames, device=device) + images_t0 = torch.zeros(n_views, 3, H, W).float().to(device) + + # Get camera viewpoints + if isinstance(elevations_deg, float) or isinstance(elevations_deg, int): + elevations_deg = [elevations_deg] * n_views + assert ( + len(elevations_deg) == n_views + ), f"Please provide 1 value, or a list of {n_views} values for elevations_deg! Given {len(elevations_deg)}" + if azimuths_deg is None: + # azimuths_deg = np.linspace(0, 360, n_views + 1)[1:] % 360 + azimuths_deg = ( + np.array([0, 60, 120, 180, 240]) + if sv4d2_model == "sv4d2" + else np.array([0, 30, 75, 120, 165, 210, 255, 300, 330]) + ) + assert ( + len(azimuths_deg) == n_views + ), f"Please provide a list of {n_views} values for azimuths_deg! Given {len(azimuths_deg)}" + polars_rad = np.array([np.deg2rad(90 - e) for e in elevations_deg]) + azimuths_rad = np.array( + [np.deg2rad((a - azimuths_deg[-1]) % 360) for a in azimuths_deg] + ) + + # Initialize image matrix + img_matrix = [[None] * n_views for _ in range(n_frames)] + for i, v in enumerate(subsampled_views): + img_matrix[0][i] = images_t0[v].unsqueeze(0) + for t in range(n_frames): + img_matrix[t][0] = images_v0[t] + + # Load SV4D++ model + model, _ = load_model( + model_config, + device, + version_dict["T"], + num_steps, + verbose, + model_path, + ) + model.en_and_decode_n_samples_a_time = decoding_t + for emb in model.conditioner.embedders: + if isinstance(emb, VideoPredictionEmbedderWithEncoder): + emb.en_and_decode_n_samples_a_time = encoding_t + + # Sampling novel-view videos + v0 = 0 + view_indices = np.arange(V) + 1 + t0_list = ( + range(0, n_frames, T) + if sv4d2_model == "sv4d2" + else range(0, n_frames - T + 1, T - 1) + ) + for t0 in tqdm(t0_list): + if t0 + T > n_frames: + t0 = n_frames - T + frame_indices = t0 + np.arange(T) + print(f"Sampling frames {frame_indices}") + image = img_matrix[t0][v0] + cond_motion = torch.cat([img_matrix[t][v0] for t in frame_indices], 0) + cond_view = torch.cat([img_matrix[t0][v] for v in view_indices], 0) + polars = polars_rad[subsampled_views[1:]][None].repeat(T, 0).flatten() + azims = azimuths_rad[subsampled_views[1:]][None].repeat(T, 0).flatten() + polars = (polars - polars_rad[v0] + torch.pi / 2) % (torch.pi * 2) + azims = (azims - azimuths_rad[v0]) % (torch.pi * 2) + cond_mv = False if t0 == 0 else True + samples = run_img2vid( + version_dict, + model, + image, + seed, + polars, + azims, + cond_motion, + cond_view, + decoding_t, + cond_mv=cond_mv, + ) + samples = samples.view(T, V, 3, H, W) + + for i, t in enumerate(frame_indices): + for j, v in enumerate(view_indices): + img_matrix[t][v] = samples[i, j][None] * 2 - 1 + + # Save output videos + for v in view_indices: + vid_file = os.path.join(output_folder, f"{base_count:06d}_v{v:03d}.mp4") + print(f"Saving {vid_file}") + save_video( + vid_file, + [img_matrix[t][v] for t in range(n_frames) if img_matrix[t][v] is not None], + ) + + +if __name__ == "__main__": + Fire(sample) diff --git a/sgm/modules/diffusionmodules/openaimodel.py b/sgm/modules/diffusionmodules/openaimodel.py index 275f4a012..e3949dd1e 100644 --- a/sgm/modules/diffusionmodules/openaimodel.py +++ b/sgm/modules/diffusionmodules/openaimodel.py @@ -74,6 +74,7 @@ def forward( x: th.Tensor, emb: th.Tensor, context: Optional[th.Tensor] = None, + cam: Optional[th.Tensor] = None, image_only_indicator: Optional[th.Tensor] = None, cond_view: Optional[th.Tensor] = None, cond_motion: Optional[th.Tensor] = None, @@ -86,7 +87,7 @@ def forward( from ...modules.spacetime_attention import ( BasicTransformerTimeMixBlock, PostHocSpatialTransformerWithTimeMixing, - PostHocSpatialTransformerWithTimeMixingAndMotion + PostHocSpatialTransformerWithTimeMixingAndMotion, ) for layer in self: @@ -97,13 +98,30 @@ def forward( ( BasicTransformerTimeMixBlock, PostHocSpatialTransformerWithTimeMixing, - PostHocSpatialTransformerWithTimeMixingAndMotion ), ): x = layer( x, context, - # cam, + emb, + time_context, + num_video_frames, + image_only_indicator, + cond_view, + cond_motion, + time_step, + name, + ) + elif isinstance( + module, + ( + PostHocSpatialTransformerWithTimeMixingAndMotion, + ), + ): + x = layer( + x, + context, + emb, time_context, num_video_frames, image_only_indicator, diff --git a/sgm/modules/diffusionmodules/video_model.py b/sgm/modules/diffusionmodules/video_model.py index 64b9844a6..9cee06657 100644 --- a/sgm/modules/diffusionmodules/video_model.py +++ b/sgm/modules/diffusionmodules/video_model.py @@ -8,10 +8,10 @@ from ...modules.spacetime_attention import ( BasicTransformerTimeMixBlock, PostHocSpatialTransformerWithTimeMixing, - PostHocSpatialTransformerWithTimeMixingAndMotion + PostHocSpatialTransformerWithTimeMixingAndMotion, ) from ...util import default -from .util import AlphaBlender # , LegacyAlphaBlenderWithBug, get_alpha +from .util import AlphaBlender, get_alpha class VideoResBlock(ResBlock): @@ -716,11 +716,11 @@ def forward( ) if self.time_mix_legacy: - alpha = self.get_alpha_fn(image_only_indicator=image_only_indicator) + alpha = self.get_alpha_fn(image_only_indicator=image_only_indicator*0.0) x = alpha.to(x.dtype) * x + (1.0 - alpha).to(x.dtype) * x_mix else: x = self.time_mixer( - x_spatial=x_mix, x_temporal=x, image_only_indicator=image_only_indicator + x_spatial=x_mix, x_temporal=x, image_only_indicator=image_only_indicator*0.0 ) x = rearrange(x, "b c t h w -> (b t) c h w") return x @@ -752,10 +752,14 @@ def __init__( context_dim: Optional[int] = None, time_downup: bool = False, time_context_dim: Optional[int] = None, + view_context_dim: Optional[int] = None, + motion_context_dim: Optional[int] = None, extra_ff_mix_layer: bool = False, use_spatial_context: bool = False, time_block_merge_strategy: str = "fixed", time_block_merge_factor: float = 0.5, + view_block_merge_factor: float = 0.5, + motion_block_merge_factor: float = 0.5, spatial_transformer_attn_type: str = "softmax", time_kernel_size: Union[int, List[int]] = 3, use_linear_in_transformer: bool = False, @@ -767,6 +771,9 @@ def __init__( max_ddpm_temb_period: int = 10000, replicate_time_mix_bug: bool = False, use_motion_attention: bool = False, + use_camera_emb: bool = False, + use_3d_attention: bool = False, + separate_motion_merge_factor: bool = False, ): super().__init__() @@ -886,11 +893,17 @@ def get_attention_layer( depth=depth, context_dim=context_dim, time_context_dim=time_context_dim, + motion_context_dim=motion_context_dim, dropout=dropout, ff_in=extra_ff_mix_layer, use_spatial_context=use_spatial_context, + use_camera_emb=use_camera_emb, + use_3d_attention=use_3d_attention, + separate_motion_merge_factor=separate_motion_merge_factor, + adm_in_channels=adm_in_channels, merge_strategy=time_block_merge_strategy, merge_factor=time_block_merge_factor, + merge_factor_motion=motion_block_merge_factor, checkpoint=use_checkpoint, use_linear=use_linear_in_transformer, attn_mode=spatial_transformer_attn_type, @@ -899,7 +912,7 @@ def get_attention_layer( time_mix_legacy=time_mix_legacy, max_time_embed_period=max_ddpm_temb_period, ) - + else: return PostHocSpatialTransformerWithTimeMixing( ch, @@ -1173,7 +1186,7 @@ def forward( timesteps: th.Tensor, context: Optional[th.Tensor] = None, y: Optional[th.Tensor] = None, - # cam: Optional[th.Tensor] = None, + cam: Optional[th.Tensor] = None, time_context: Optional[th.Tensor] = None, num_video_frames: Optional[int] = None, image_only_indicator: Optional[th.Tensor] = None, @@ -1199,7 +1212,7 @@ def forward( h, emb, context=context, - # cam=cam, + cam=cam, image_only_indicator=image_only_indicator, cond_view=cond_view, cond_motion=cond_motion, @@ -1213,7 +1226,7 @@ def forward( h, emb, context=context, - # cam=cam, + cam=cam, image_only_indicator=image_only_indicator, cond_view=cond_view, cond_motion=cond_motion, @@ -1228,7 +1241,7 @@ def forward( h, emb, context=context, - # cam=cam, + cam=cam, image_only_indicator=image_only_indicator, cond_view=cond_view, cond_motion=cond_motion, diff --git a/sgm/modules/spacetime_attention.py b/sgm/modules/spacetime_attention.py index 1ca44bb2d..370879588 100644 --- a/sgm/modules/spacetime_attention.py +++ b/sgm/modules/spacetime_attention.py @@ -1,6 +1,7 @@ from functools import partial import torch +import torch.nn.functional as F from ..modules.attention import * from ..modules.diffusionmodules.util import ( @@ -359,11 +360,17 @@ def __init__( use_linear=False, context_dim=None, use_spatial_context=False, + use_camera_emb=False, + use_3d_attention=False, + separate_motion_merge_factor=False, + adm_in_channels=None, timesteps=None, merge_strategy: str = "fixed", merge_factor: float = 0.5, + merge_factor_motion: float = 0.5, apply_sigmoid_to_merge: bool = True, time_context_dim=None, + motion_context_dim=None, ff_in=False, checkpoint=False, time_depth=1, @@ -388,6 +395,10 @@ def __init__( self.time_depth = time_depth self.depth = depth self.max_time_embed_period = max_time_embed_period + self.use_camera_emb = use_camera_emb + self.motion_context_dim = motion_context_dim + self.use_3d_attention = use_3d_attention + self.separate_motion_merge_factor = separate_motion_merge_factor time_mix_d_head = d_head n_time_mix_heads = n_heads @@ -398,9 +409,6 @@ def __init__( if use_spatial_context: time_context_dim = context_dim - camera_context_dim = time_context_dim - motion_context_dim = 4 # time_context_dim - # Camera attention layer self.time_mix_blocks = nn.ModuleList( [ @@ -409,7 +417,7 @@ def __init__( n_time_mix_heads, time_mix_d_head, dropout=dropout, - context_dim=camera_context_dim, + context_dim=time_context_dim, timesteps=timesteps, checkpoint=checkpoint, ff_in=ff_in, @@ -449,9 +457,10 @@ def __init__( self.in_channels = in_channels time_embed_dim = self.in_channels * 4 + time_embed_channels = adm_in_channels if self.use_camera_emb else self.in_channels # Camera view embedding self.time_mix_time_embed = nn.Sequential( - linear(self.in_channels, time_embed_dim), + linear(time_embed_channels, time_embed_dim), nn.SiLU(), linear(time_embed_dim, self.in_channels), ) @@ -486,12 +495,16 @@ def __init__( self.time_mixer = AlphaBlender( alpha=merge_factor, merge_strategy=merge_strategy ) + if self.separate_motion_merge_factor: + self.time_mixer_motion = AlphaBlender( + alpha=merge_factor_motion, merge_strategy=merge_strategy + ) def forward( self, x: torch.Tensor, context: Optional[torch.Tensor] = None, - # cam: Optional[torch.Tensor] = None, + cam: Optional[torch.Tensor] = None, time_context: Optional[torch.Tensor] = None, timesteps: Optional[int] = None, image_only_indicator: Optional[torch.Tensor] = None, @@ -500,16 +513,19 @@ def forward( time_step: Optional[int] = None, name: Optional[str] = None, ) -> torch.Tensor: - _, _, h, w = x.shape + # context: b t 1024 + # cond_view: b*v 4 h w + # cond_motion: b*t 4 h w + # image_only_indicator: b t*v + b, t, d1 = context.shape # CLIP + v, d2 = cond_view.shape[0]//b, cond_view.shape[1] # VAE + _, c, h, w = x.shape + x_in = x spatial_context = None if exists(context): spatial_context = context - # cond_view: b v 4 h w - # cond_motion: b t 4 h w - b, t, d1 = context.shape # CLIP - v, d2 = cond_view.shape[0]//b, cond_view.shape[1] # VAE cond_view = torch.nn.functional.interpolate(cond_view, size=(h,w), mode="bilinear") # b*v d h w spatial_context = context[:,:,None].repeat(1,1,v,1).reshape(-1,1,d1) # (b*t*v) 1 d1 camera_context = context[:,:,None].repeat(1,1,h*w,1).reshape(-1,1,d1) # (b*t*h*w) 1 d1 @@ -518,10 +534,9 @@ def forward( x = self.norm(x) if not self.use_linear: x = self.proj_in(x) - x = rearrange(x, "b c h w -> b (h w) c") # 21 x 4096 x 320 + x = rearrange(x, "b c h w -> b (h w) c") if self.use_linear: x = self.proj_in(x) - c = x.shape[-1] if self.time_mix_legacy: alpha = self.get_alpha_fn(image_only_indicator=image_only_indicator) @@ -536,19 +551,26 @@ def forward( max_period=self.max_time_embed_period, ) emb_time = self.time_mix_motion_embed(t_emb) - emb_time = emb_time[:, None, :] # b*t x 1 x 320 + emb_time = emb_time[:, None, :] # b*t 1 c - num_views = torch.arange(v, device=x.device) - num_views = repeat(num_views, "t -> b t", b=b) - num_views = rearrange(num_views, "b t -> (b t)") - v_emb = timestep_embedding( - num_views, - self.in_channels, - repeat_only=False, - max_period=self.max_time_embed_period, - ) - emb_view = self.time_mix_time_embed(v_emb) - emb_view = emb_view[:, None, :] # b*v x 1 x 320 + if self.use_camera_emb: + emb_view = self.time_mix_time_embed(cam.view(b,t,v,-1)[:,0].reshape(b*v,-1)) + emb_view = emb_view[:, None, :] + else: + num_views = torch.arange(v, device=x.device) + num_views = repeat(num_views, "t -> b t", b=b) + num_views = rearrange(num_views, "b t -> (b t)") + v_emb = timestep_embedding( + num_views, + self.in_channels, + repeat_only=False, + max_period=self.max_time_embed_period, + ) + emb_view = self.time_mix_time_embed(v_emb) + emb_view = emb_view[:, None, :] # b*v 1 c + + if self.use_3d_attention: + emb_view = emb_view.repeat(1, h*w, 1).view(-1,1,c) # b*v*h*w 1 c for it_, (block, time_block, mot_block) in enumerate( zip(self.transformer_blocks, self.time_mix_blocks, self.motion_blocks) @@ -560,7 +582,10 @@ def forward( ) # Camera attention - x = x.view(b, t, v, h*w, c).permute(0,2,1,3,4).reshape(b*v,-1,c) # b*v t*h*w c + if self.use_3d_attention: + x = x.view(b, t, v, h*w, c).permute(0,2,3,1,4).reshape(-1,t,c) # b*v*h*w t c + else: + x = x.view(b, t, v, h*w, c).permute(0,2,1,3,4).reshape(b*v,-1,c) # b*v t*h*w c x_mix = x + emb_view x_mix = time_block(x_mix, context=camera_context, timesteps=v) if self.time_mix_legacy: @@ -569,20 +594,24 @@ def forward( x = self.time_mixer( x_spatial=x, x_temporal=x_mix, - image_only_indicator=image_only_indicator[:,:v], + image_only_indicator=torch.zeros_like(image_only_indicator[:,:1].repeat(1,x.shape[0]//b)), ) # Motion attention - x = x.view(b, v, t, h*w, c).permute(0,2,1,3,4).reshape(b*t,-1,c) # b*t v*h*w c + if self.use_3d_attention: + x = x.view(b, v, h*w, t, c).permute(0,3,1,2,4).reshape(b*t,-1,c) # b*t v*h*w c + else: + x = x.view(b, v, t, h*w, c).permute(0,2,1,3,4).reshape(b*t,-1,c) # b*t v*h*w c x_mix = x + emb_time x_mix = mot_block(x_mix, context=motion_context, timesteps=t) if self.time_mix_legacy: x = alpha.to(x.dtype) * x + (1.0 - alpha).to(x.dtype) * x_mix else: - x = self.time_mixer( + motion_mixer = self.time_mixer_motion if self.separate_motion_merge_factor else self.time_mixer + x = motion_mixer( x_spatial=x, x_temporal=x_mix, - image_only_indicator=image_only_indicator[:,:t], + image_only_indicator=torch.zeros_like(image_only_indicator[:,:1].repeat(1,x.shape[0]//b)), ) x = x.view(b, t, v, h*w, c).reshape(-1,h*w,c) # b*t*v h*w c