diff --git a/scripts/demo/sv4d_helpers.py b/scripts/demo/sv4d_helpers.py index d67231d..96a5d0b 100755 --- a/scripts/demo/sv4d_helpers.py +++ b/scripts/demo/sv4d_helpers.py @@ -121,10 +121,6 @@ def save_video(file_name, imgs, fps=10): def read_video( input_path: str, n_frames: int, - W: int, - H: int, - remove_bg: bool = False, - image_frame_ratio: Optional[float] = None, device: str = "cuda", ): path = Path(input_path) @@ -158,47 +154,121 @@ def read_video( if len(images) < n_frames: images = (images + images[::-1])[:n_frames] - if len(images) != n_frames: raise ValueError(f"Input video contains fewer than {n_frames} frames.") - # Remove background and crop video frames images_v0 = [] - for t, image in enumerate(images): - if remove_bg: - if image.mode != "RGBA": - image.thumbnail([W, H], Image.Resampling.LANCZOS) - image = remove(image.convert("RGBA"), alpha_matting=True) - image_arr = np.array(image) - in_w, in_h = image_arr.shape[:2] - ret, mask = cv2.threshold( - np.array(image.split()[-1]), 0, 255, cv2.THRESH_BINARY - ) - x, y, w, h = cv2.boundingRect(mask) - max_size = max(w, h) - if t == 0: - side_len = ( - int(max_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 - padded_image[ - center - h // 2 : center - h // 2 + h, - center - w // 2 : center - w // 2 + w, - ] = image_arr[y : y + h, x : x + w] - rgba = Image.fromarray(padded_image).resize((W, H), Image.LANCZOS) - rgba_arr = np.array(rgba) / 255.0 - rgb = rgba_arr[..., :3] * rgba_arr[..., -1:] + (1 - rgba_arr[..., -1:]) - image = Image.fromarray((rgb * 255).astype(np.uint8)) - else: - image = image.convert("RGB").resize((W, H), Image.LANCZOS) + + for image in images: image = ToTensor()(image).unsqueeze(0).to(device) images_v0.append(image * 2.0 - 1.0) 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): + print(f"preprocess {input_path}") + if output_folder is None: + output_folder = os.path.dirname(input_path) + path = Path(input_path) + is_video_file = False + all_img_paths = [] + if path.is_file(): + if any([input_path.endswith(x) for x in [".gif", ".mp4"]]): + is_video_file = True + else: + raise ValueError("Path is not a valid video file.") + elif path.is_dir(): + all_img_paths = sorted( + [ + f + for f in path.iterdir() + if f.is_file() and f.suffix.lower() in [".jpg", ".jpeg", ".png"] + ] + )[:n_frames] + elif "*" in input_path: + all_img_paths = sorted(glob(input_path))[:n_frames] + else: + raise ValueError + + if is_video_file and input_path.endswith(".gif"): + images = read_gif(input_path, n_frames)[:n_frames] + elif is_video_file and input_path.endswith(".mp4"): + images = read_mp4(input_path, n_frames)[:n_frames] + else: + print(f"Loading {len(all_img_paths)} video frames...") + 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.") + + # Remove background + for i, image in enumerate(images): + if remove_bg: + if image.mode == "RGBA": + pass + else: + # image.thumbnail([W, H], Image.Resampling.LANCZOS) + image = remove(image.convert("RGBA"), alpha_matting=True) + images[i] = image + + # Crop video frames, assume the object is already in the center of the image + white_thresh = 250 + images_v0 = [] + box_coord = [np.inf, np.inf, 0, 0] + for image in images: + image_arr = np.array(image) + in_w, in_h = image_arr.shape[:2] + original_center = (in_w // 2, in_h // 2) + if image.mode == "RGBA": + ret, mask = cv2.threshold( + np.array(image.split()[-1]), 0, 255, cv2.THRESH_BINARY + ) + 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 + ) + + 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(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 + box_size = box_square * 2 + + for image in images: + if image.mode == "RGB": + 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 + ) + padded_image = np.zeros((side_len, side_len, 4), dtype=np.uint8) + center = side_len // 2 + padded_image[ + center - box_size // 2 : center - box_size // 2 + box_size, + center - box_size // 2 : center - box_size // 2 + box_size, + ] = image_arr[x : x + w, y : y + h] + + rgba = Image.fromarray(padded_image).resize((W, H), Image.LANCZOS) + # rgba = image.resize((W, H), Image.LANCZOS) + 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"))) // 10 + 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 @@ -212,26 +282,32 @@ def sample_sv3d( polar_rad: Optional[Union[float, List[float]]] = None, azim_rad: Optional[List[float]] = None, verbose: Optional[bool] = False, + sv3d_model=None, ): """ Simple script to generate a single sample conditioned on an image `input_path` or multiple images, one for each image file in folder `input_path`. If you run out of VRAM, try decreasing `decoding_t`. """ - if version == "sv3d_u": - model_config = "scripts/sampling/configs/sv3d_u.yaml" - elif version == "sv3d_p": - model_config = "scripts/sampling/configs/sv3d_p.yaml" - else: - raise ValueError(f"Version {version} does not exist.") + if sv3d_model is None: + if version == "sv3d_u": + model_config = "scripts/sampling/configs/sv3d_u.yaml" + elif version == "sv3d_p": + model_config = "scripts/sampling/configs/sv3d_p.yaml" + else: + raise ValueError(f"Version {version} does not exist.") - model, filter = load_model( - model_config, - device, - num_frames, - num_steps, - verbose, - ) + model, filter = load_model( + model_config, + device, + num_frames, + num_steps, + verbose, + ) + else: + model = sv3d_model + + load_module_gpu(model) H, W = image.shape[2:] F = 8 @@ -286,25 +362,32 @@ def sample_sv3d( ) samples_z = model.sampler(denoiser, randn, cond=c, uc=uc) + unload_module_gpu(model.model) + unload_module_gpu(model.denoiser) model.en_and_decode_n_samples_a_time = decoding_t samples_x = model.decode_first_stage(samples_z) samples_x[-1:] = value_dict["cond_frames_without_noise"] samples = torch.clamp(samples_x, min=-1.0, max=1.0) - return samples - - -def decode_latents(model, samples_z, timesteps): - load_module_gpu(model.first_stage_model) - if isinstance(model.first_stage_model.decoder, VideoDecoder): - samples_x = model.decode_first_stage(samples_z, timesteps=timesteps) - else: - 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) + unload_module_gpu(model) return samples +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 isinstance(model.first_stage_model.decoder, VideoDecoder): + 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) + img_matrix[t][v] = samples * 2 - 1 + unload_module_gpu(model.first_stage_model) + return img_matrix + + def init_embedder_options_no_st(keys, init_dict, prompt=None, negative_prompt=None): # Hardcoded demo settings; might undergo some changes in the future @@ -604,6 +687,7 @@ def run_img2vid( azim_rad=np.linspace(0, 360, 21 + 1)[1:], cond_motion=None, cond_view=None, + decoding_t=None, ): options = version_dict["options"] H = version_dict["H"] @@ -670,12 +754,53 @@ def run_img2vid( force_uc_zero_embeddings=options.get("force_uc_zero_embeddings", None), force_cond_zero_embeddings=options.get("force_cond_zero_embeddings", None), return_latents=False, - decoding_t=options.get("decoding_T", T), + decoding_t=decoding_t, ) return samples +def prepare_inputs(frame_indices, img_matrix, v0, view_indices, model, version_dict, seed, polars, azims): + load_module_gpu(model.conditioner) + + 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, + ) + + # backward sampling + 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, + ) + + unload_module_gpu(model.conditioner) + return forward_inputs, forward_frame_indices, backward_inputs, backward_frame_indices + def do_sample( model, sampler, @@ -722,6 +847,8 @@ def do_sample( force_cond_zero_embeddings=force_cond_zero_embeddings, ) unload_module_gpu(model.conditioner) + print("anchor_after_condition {}".format(torch.cuda.memory_reserved() / (1024 ** 3))) + # torch.cuda.empty_cache() for k in c: if not k == "crossattn": @@ -761,14 +888,15 @@ def do_sample( 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) - + print("anchor_after_denoiser {}".format(torch.cuda.memory_reserved() / (1024 ** 3))) + # torch.cuda.empty_cache() load_module_gpu(model.first_stage_model) + model.en_and_decode_n_samples_a_time = decoding_t if isinstance(model.first_stage_model.decoder, VideoDecoder): samples_x = model.decode_first_stage( samples_z, timesteps=default(decoding_t, T) @@ -777,17 +905,16 @@ def do_sample( 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 - + # torch.cuda.empty_cache() return samples -def do_sample_per_step( +def prepare_sampling_( model, sampler, value_dict, @@ -797,8 +924,6 @@ def do_sample_per_step( batch2model_input: List = None, T=None, additional_batch_uc_fields=None, - step=None, - noisy_latents=None, ): force_uc_zero_embeddings = default(force_uc_zero_embeddings, []) batch2model_input = default(batch2model_input, []) @@ -812,8 +937,6 @@ def do_sample_per_step( 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, @@ -827,8 +950,7 @@ def do_sample_per_step( force_uc_zero_embeddings=force_uc_zero_embeddings, force_cond_zero_embeddings=force_cond_zero_embeddings, ) - unload_module_gpu(model.conditioner) - + print("dense_after_condition {}".format(torch.cuda.memory_reserved() / (1024 ** 3))) for k in c: if not k == "crossattn": c[k], uc[k] = map( @@ -859,7 +981,14 @@ def do_sample_per_step( ) 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): + precision_scope = autocast + with torch.no_grad(): + with precision_scope("cuda"): + with model.ema_scope(): noisy_latents_scaled, s_in, sigmas, num_sigmas, _, _ = ( sampler.prepare_sampling_loop( noisy_latents.clone(), c, uc, sampler.num_steps @@ -893,13 +1022,11 @@ def do_sample_per_step( uc, gamma, ) - unload_module_gpu(model.model) - unload_module_gpu(model.denoiser) - + print("dense_after_sampling {}".format(torch.cuda.memory_reserved() / (1024 ** 3))) return samples_z -def run_img2vid_per_step( +def prepare_sampling( version_dict, model, image, @@ -908,8 +1035,6 @@ def run_img2vid_per_step( azim_rad=np.linspace(0, 360, 21 + 1)[1:], cond_motion=None, cond_view=None, - step=None, - noisy_latents=None, ): options = version_dict["options"] H = version_dict["H"] @@ -962,7 +1087,7 @@ def run_img2vid_per_step( sampler, num_rows, num_cols = init_sampling_no_st(options=options) num_samples = num_rows * num_cols - samples = do_sample_per_step( + c, uc, additional_model_inputs = prepare_sampling_( model, sampler, value_dict, @@ -971,11 +1096,9 @@ def run_img2vid_per_step( force_cond_zero_embeddings=options.get("force_cond_zero_embeddings", None), batch2model_input=["num_video_frames", "image_only_indicator"], T=T, - step=step, - noisy_latents=noisy_latents, ) - return samples + return c, uc, additional_model_inputs, sampler def get_unique_embedder_keys_from_conditioner(conditioner): diff --git a/scripts/sampling/configs/sv4d.yaml b/scripts/sampling/configs/sv4d.yaml index 56e6ddb..ded3f23 100755 --- a/scripts/sampling/configs/sv4d.yaml +++ b/scripts/sampling/configs/sv4d.yaml @@ -1,6 +1,7 @@ N_TIME: 5 N_VIEW: 8 N_FRAMES: 40 +ENCODE_N_A_TIME: 8 model: target: sgm.models.diffusion.DiffusionEngine @@ -67,6 +68,7 @@ model: is_ae: True n_cond_frames: ${N_FRAMES} n_copies: 1 + en_and_decode_n_samples_a_time: ${ENCODE_N_A_TIME} encoder_config: target: sgm.models.autoencoder.AutoencoderKLModeOnly params: @@ -131,6 +133,7 @@ model: is_ae: True n_cond_frames: ${N_VIEW} n_copies: 1 + en_and_decode_n_samples_a_time: ${ENCODE_N_A_TIME} sigma_sampler_config: target: sgm.modules.diffusionmodules.sigma_sampling.ZeroSampler @@ -141,6 +144,7 @@ model: is_ae: True n_cond_frames: ${N_TIME} n_copies: 1 + en_and_decode_n_samples_a_time: ${ENCODE_N_A_TIME} encoder_config: target: sgm.models.autoencoder.AutoencoderKLModeOnly params: diff --git a/scripts/sampling/simple_video_sample_4d.py b/scripts/sampling/simple_video_sample_4d.py index c970e74..a318069 100755 --- a/scripts/sampling/simple_video_sample_4d.py +++ b/scripts/sampling/simple_video_sample_4d.py @@ -16,9 +16,12 @@ from scripts.demo.sv4d_helpers import ( initial_model_load, read_video, run_img2vid, - run_img2vid_per_step, + prepare_sampling, + prepare_inputs, + do_sample_per_step, sample_sv3d, save_video, + preprocess_video, ) @@ -32,11 +35,11 @@ def sample( motion_bucket_id: int = 127, cond_aug: float = 1e-5, seed: int = 23, - decoding_t: int = 14, # Number of frames decoded 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[Union[float, List[float]]] = 10.0, azimuths_deg: Optional[List[float]] = None, - image_frame_ratio: Optional[float] = None, + image_frame_ratio: Optional[float] = 0.917, verbose: Optional[bool] = False, remove_bg: bool = False, ): @@ -89,15 +92,16 @@ def sample( # Read input video frames i.e. images at view 0 print(f"Reading {input_path}") - images_v0 = read_video( + processed_input_path = preprocess_video( input_path, + remove_bg=remove_bg, n_frames=n_frames, W=W, H=H, - remove_bg=remove_bg, + output_folder=output_folder, image_frame_ratio=image_frame_ratio, - device=device, ) + images_v0 = read_video(processed_input_path, n_frames=n_frames, device=device) # Get camera viewpoints if isinstance(elevations_deg, float) or isinstance(elevations_deg, int): @@ -139,7 +143,7 @@ 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"))) // 11 + 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], @@ -171,7 +175,7 @@ def sample( azims = azimuths_rad[subsampled_views[1:]][None].repeat(T, 0).flatten() azims = (azims - azimuths_rad[v0]) % (torch.pi * 2) samples = run_img2vid( - version_dict, model, image, seed, polars, azims, cond_motion, cond_view + version_dict, model, image, seed, polars, azims, cond_motion, cond_view, decoding_t ) samples = samples.view(T, V, 3, H, W) for i, t in enumerate(frame_indices): @@ -185,40 +189,48 @@ def sample( frame_indices = t0 + np.arange(T) print(f"Sampling dense frames {frame_indices}") latent_matrix = torch.randn(n_frames, n_views, C, H // F, W // F).to("cuda") + + polars = polars_rad[subsampled_views[1:]][None].repeat(T, 0).flatten() + azims = azimuths_rad[subsampled_views[1:]][None].repeat(T, 0).flatten() + azims = (azims - azimuths_rad[v0]) % (torch.pi * 2) + + # alternate between forward and backward conditioning + forward_inputs, forward_frame_indices, backward_inputs, backward_frame_indices = prepare_inputs( + frame_indices, + img_matrix, + v0, + view_indices, + model, + version_dict, + seed, + polars, + azims + ) + for step in tqdm(range(num_steps)): - frame_indices = frame_indices[ - ::-1 - ].copy() # alternate between forward and backward conditioning - t0 = frame_indices[0] - 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() - azims = (azims - azimuths_rad[v0]) % (torch.pi * 2) + if step % 2 == 1: + c, uc, additional_model_inputs, sampler = forward_inputs + frame_indices = forward_frame_indices + else: + c, uc, additional_model_inputs, sampler = backward_inputs + frame_indices = backward_frame_indices noisy_latents = latent_matrix[frame_indices][:, view_indices].flatten(0, 1) - samples = run_img2vid_per_step( - version_dict, + + samples = do_sample_per_step( model, - image, - seed, - polars, - azims, - cond_motion, - cond_view, - step, + sampler, noisy_latents, + c, + uc, + step, + additional_model_inputs, ) samples = samples.view(T, V, C, H // F, W // F) for i, t in enumerate(frame_indices): for j, v in enumerate(view_indices): latent_matrix[t, v] = samples[i, j] - for t in frame_indices: - for v in view_indices: - if t != 0 and v != 0: - img = decode_latents(model, latent_matrix[t, v][None], T) - img_matrix[t][v] = img * 2 - 1 + img_matrix = decode_latents(model, latent_matrix, img_matrix, frame_indices, view_indices, T) # Save output videos for v in view_indices: