commit f2ece13a51a07e1cc35930ec18ead742ef71e2c1 Author: AIWarper Date: Fri May 9 12:18:59 2025 -0600 Initial commit of ComfyUI-NormalCrafter node diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e298a8e --- /dev/null +++ b/.gitignore @@ -0,0 +1,42 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# Virtualenv +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# IDE / Editor specific +.vscode/ +.idea/ +*.swp +*.swo + +# OS specific +.DS_Store +Thumbs.db \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..6fb3108 --- /dev/null +++ b/README.md @@ -0,0 +1,56 @@ +# ComfyUI NormalCrafter + +This is a ComfyUI custom node implementation for [NormalCrafter: Learning Temporally Consistent Normals from Video Diffusion Priors](https://github.com/Binyr/NormalCrafter) by Yanrui Bin, Wenbo Hu, Haoyuan Wang, Xinya Chen, and Bing Wang. + +It allows you to generate temporally consistent normal map sequences from input video frames. + +## Installation + +1. Clone this repository into your `ComfyUI/custom_nodes/` directory: + ```bash + cd ComfyUI/custom_nodes/ + git clone https://github.com/YOUR_USERNAME/ComfyUI-NormalCrafter.git # Replace with your actual repo URL + ``` +2. Install the required dependencies. Navigate to the `ComfyUI-NormalCrafter` directory and install using the `requirements.txt`: + ```bash + cd ComfyUI-NormalCrafter + pip install -r requirements.txt + ``` + (See the `requirements.txt` for notes on the `diffusers` dependency, which ComfyUI often manages.) +3. Restart ComfyUI. + +## Node: NormalCrafter (Process Video) + +This node takes a sequence of images (video frames) and processes them to output a corresponding sequence of normal map images. + +### Parameters + +* **`images` (Input Socket)**: The input image sequence (video frames). +* **`pipe_override` (Input Socket, Optional)**: Allows providing a pre-loaded NormalCrafter pipeline instance. If unconnected, the node loads its own. +* **`seed`**: (Integer, Default: 42) Controls the randomness for reproducible results. +* **`control_after_generate`**: (Fixed, Increment, Decrement, Randomize) Standard ComfyUI widget for seed behavior on subsequent runs. Note: The underlying pipeline uses the seed for each full video processing. +* **`max_res_dimension`**: (Integer, Default: 1024) The maximum dimension (height or width) to which input frames are resized while maintaining aspect ratio. +* **`window_size`**: (Integer, Default: 14) The number of consecutive frames processed together in a sliding window. Affects temporal consistency. +* **`time_step_size`**: (Integer, Default: 10) How many frames the sliding window moves forward after processing a chunk. If less than `window_size`, frames will overlap, potentially improving smoothness. +* **`decode_chunk_size`**: (Integer, Default: 4) Number of latent frames decoded by the VAE at once. Primarily a VRAM management setting. +* **`fps_for_time_ids`**: (Integer, Default: 7) Conditions the model on an intended Frames Per Second, influencing motion characteristics in the generated normals. +* **`motion_bucket_id`**: (Integer, Default: 127) Conditions the model on an expected amount of motion. +* **`noise_aug_strength`**: (Float, Default: 0.0) Strength of noise augmentation applied to conditioning information. *Note: In testing, this parameter showed minimal to no visible effect on the output for this specific model and task.* + +### Troubleshooting Flicker / Improving Temporal Consistency + +If you are experiencing flickering or temporal inconsistencies in your output: + +* **Increase `window_size`**: A larger window allows the model to see more temporal context, which can significantly improve consistency between frames. +* **Adjust `time_step_size`**: Using a `time_step_size` smaller than `window_size` creates an overlap between processed windows. This overlap is merged, which can smooth transitions. For example, if `window_size` is 20, try a `time_step_size` of 10 or 15. + +You may be able to increase `window_size` and `time_step_size` substantially (e.g., to their maximum values) without encountering Out Of Memory (OOM) issues, depending on your hardware. Experiment to find the best balance for your needs. + +### Dependencies + +* `mediapy` +* `decord` +* `diffusers` (and its dependencies like `transformers`, `huggingface_hub`) - ComfyUI usually manages its own `diffusers` version. Install manually if you encounter specific import errors related to it. +* `torch`, `numpy`, `Pillow` (standard Python ML/Image libraries) + +Refer to `requirements.txt` for more details. \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3205a0d --- /dev/null +++ b/__init__.py @@ -0,0 +1,21 @@ +# ComfyUI/custom_nodes/ComfyUI-NormalCrafter/__init__.py + +try: + from .normal_crafter_nodes import NormalCrafterNode + + NODE_CLASS_MAPPINGS = { + "NormalCrafterNode": NormalCrafterNode, + } + + NODE_DISPLAY_NAME_MAPPINGS = { + "NormalCrafterNode": "NormalCrafter (Process Video)", + } + + __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + print("✅ ComfyUI-NormalCrafter: Custom nodes loaded successfully.") + +except ImportError as e: + print(f"❌ ComfyUI-NormalCrafter: Failed to import nodes: {e}") + print(" Please ensure all dependencies are installed (e.g., diffusers, transformers, huggingface_hub) and the files are correctly placed.") + NODE_CLASS_MAPPINGS = {} + NODE_DISPLAY_NAME_MAPPINGS = {} \ No newline at end of file diff --git a/git b/git new file mode 100644 index 0000000..e69de29 diff --git a/normal_crafter_nodes.py b/normal_crafter_nodes.py new file mode 100644 index 0000000..f1637f2 --- /dev/null +++ b/normal_crafter_nodes.py @@ -0,0 +1,246 @@ +# ComfyUI/custom_nodes/ComfyUI-NormalCrafter/normal_crafter_nodes.py + +import torch +import numpy as np +from PIL import Image +import os + +import comfy.model_management +import comfy.utils +import folder_paths # ComfyUI's way to get model paths + +# Try to import required components and provide guidance if they are missing +try: + from diffusers import AutoencoderKLTemporalDecoder + from huggingface_hub import snapshot_download +except ImportError: + print("ComfyUI-NormalCrafter: Missing essential libraries. Please ensure 'diffusers', 'transformers', 'accelerate', and 'huggingface_hub' are installed.") + # Consider raising an error or making the node unusable if critical dependencies are missing. + +try: + # This assumes the 'normalcrafter' directory is correctly placed within ComfyUI-NormalCrafter + from .normalcrafter.normal_crafter_ppl import NormalCrafterPipeline + from .normalcrafter.unet import DiffusersUNetSpatioTemporalConditionModelNormalCrafter +except ImportError as e: + print(f"ComfyUI-NormalCrafter: Error importing NormalCrafter components: {e}. " + "Ensure the 'normalcrafter' directory is correctly placed inside 'ComfyUI-NormalCrafter'.") + +# Global variable to cache the pipeline +NORMALCRAFTER_PIPE = None +CURRENT_PIPE_CONFIG = {} + +# Define model paths and repo ID +NORMALCRAFTER_REPO_ID = "Yanrui95/NormalCrafter" +NORMALCRAFTER_MODELS_SUBDIR_NAME = "normalcrafter_models" # Subdirectory in ComfyUI/models/ +SVD_XT_REPO_ID = "stabilityai/stable-video-diffusion-img2vid-xt" + + +class NormalCrafterNode: + def __init__(self): + self.pipe = None # Instance variable for the pipeline + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}), + "max_res_dimension": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 64}), + "window_size": ("INT", {"default": 14, "min": 2, "max": 64}), + "time_step_size": ("INT", {"default": 10, "min": 1, "max": 64}), + "decode_chunk_size": ("INT", {"default": 4, "min": 1, "max": 64}), + # "fps_for_time_ids": ("INT", {"default": 7, "min": 1, "max": 60}), # <<< COMMENTED OUT + # "motion_bucket_id": ("INT", {"default": 127, "min": 0, "max": 511}), # <<< COMMENTED OUT + # "noise_aug_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), # <<< COMMENTED OUT + }, + "optional": { + "pipe_override": ("NORMALCRAFTER_PIPE",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "process" + CATEGORY = "Video/NormalCrafter" + + def _get_local_nc_model_path(self): + base_models_dir = os.path.join(folder_paths.models_dir, NORMALCRAFTER_MODELS_SUBDIR_NAME) + os.makedirs(base_models_dir, exist_ok=True) + model_name_folder = NORMALCRAFTER_REPO_ID.split('/')[-1] + specific_model_path = os.path.join(base_models_dir, model_name_folder) + return specific_model_path + + def _download_normalcrafter_model_if_needed(self): + local_nc_path = self._get_local_nc_model_path() + unet_config_path = os.path.join(local_nc_path, "unet", "config.json") + + if not os.path.exists(unet_config_path): + print(f"ComfyUI-NormalCrafter: Downloading {NORMALCRAFTER_REPO_ID} model to {local_nc_path}...") + try: + snapshot_download( + repo_id=NORMALCRAFTER_REPO_ID, + local_dir=local_nc_path, + local_dir_use_symlinks=False, + ) + print(f"ComfyUI-NormalCrafter: Model {NORMALCRAFTER_REPO_ID} download complete.") + except Exception as e: + print(f"ComfyUI-NormalCrafter: Failed to download model {NORMALCRAFTER_REPO_ID}: {e}") + raise + else: + # print(f"ComfyUI-NormalCrafter: Model {NORMALCRAFTER_REPO_ID} found at {local_nc_path}.") + pass + return local_nc_path + + def _load_pipeline(self, device_str="cuda", dtype_str="float16"): + global NORMALCRAFTER_PIPE, CURRENT_PIPE_CONFIG + + requested_config = { + "device": device_str, + "dtype": dtype_str, + } + + if NORMALCRAFTER_PIPE is not None and CURRENT_PIPE_CONFIG == requested_config: + target_device = comfy.model_management.get_torch_device() if device_str == "cuda" else torch.device("cpu") + if NORMALCRAFTER_PIPE.device != target_device: + NORMALCRAFTER_PIPE.to(target_device) + self.pipe = NORMALCRAFTER_PIPE + return self.pipe + + print("ComfyUI-NormalCrafter: Loading NormalCrafter pipeline...") + local_nc_model_path = self._download_normalcrafter_model_if_needed() + + if dtype_str == "float16": torch_dtype = torch.float16 + elif dtype_str == "bf16": torch_dtype = torch.bfloat16 + else: torch_dtype = torch.float32 + + print(f"ComfyUI-NormalCrafter: Loading UNet from subfolder 'unet' in {local_nc_model_path}") + unet = DiffusersUNetSpatioTemporalConditionModelNormalCrafter.from_pretrained( + local_nc_model_path, subfolder="unet", low_cpu_mem_usage=True, torch_dtype=torch_dtype + ) + print(f"ComfyUI-NormalCrafter: Loading VAE from subfolder 'vae' in {local_nc_model_path}") + vae = AutoencoderKLTemporalDecoder.from_pretrained( + local_nc_model_path, subfolder="vae", low_cpu_mem_usage=True, torch_dtype=torch_dtype + ) + + svd_xt_variant = "fp16" if torch_dtype == torch.float16 else None + print(f"ComfyUI-NormalCrafter: Loading base SVD pipeline ({SVD_XT_REPO_ID}) and integrating custom UNet/VAE.") + pipe = NormalCrafterPipeline.from_pretrained( + SVD_XT_REPO_ID, unet=unet, vae=vae, torch_dtype=torch_dtype, variant=svd_xt_variant, + ) + + target_device = comfy.model_management.get_torch_device() if device_str == "cuda" else torch.device("cpu") + try: + pipe.enable_xformers_memory_efficient_attention() + print("ComfyUI-NormalCrafter: Xformers memory efficient attention enabled.") + except Exception as e: + print(f"ComfyUI-NormalCrafter: Xformers not available or failed to enable: {e}. Proceeding without it.") + + pipe.to(target_device) + NORMALCRAFTER_PIPE = pipe + CURRENT_PIPE_CONFIG = requested_config + self.pipe = NORMALCRAFTER_PIPE + print(f"ComfyUI-NormalCrafter: Pipeline loaded to {target_device} with dtype {torch_dtype}.") + return self.pipe + + def tensor_to_pil_list(self, images_tensor: torch.Tensor) -> list: + pil_images = [] + for i in range(images_tensor.shape[0]): + img_np = (images_tensor[i].cpu().numpy() * 255).astype(np.uint8) + pil_images.append(Image.fromarray(img_np)) + return pil_images + + def resize_pil_images(self, pil_images: list, max_res_dim: int) -> list: + resized_images = [] + if not pil_images: return [] + for img in pil_images: + original_width, original_height = img.size + if max(original_height, original_width) > max_res_dim: + scale = max_res_dim / max(original_height, original_width) + target_height = round(original_height * scale) + target_width = round(original_width * scale) + else: + target_height = original_height + target_width = original_width + resized_images.append(img.resize((target_width, target_height), Image.LANCZOS)) + return resized_images + + def process(self, images: torch.Tensor, seed: int, max_res_dimension: int, + window_size: int, time_step_size: int, decode_chunk_size: int, + # fps_for_time_ids: int, motion_bucket_id: int, noise_aug_strength: float, # <<< COMMENTED OUT from signature + pipe_override=None): + + # --- Define default values for the parameters we removed from UI --- + default_fps_for_time_ids = 7 + default_motion_bucket_id = 127 + default_noise_aug_strength = 0.0 + # --- + + # The DEBUG ENTRY print statement that used these parameters would be removed or modified + # print(f"DEBUG ENTRY: seed={seed}, fps={default_fps_for_time_ids}, motion_id={default_motion_bucket_id}, noise_aug={default_noise_aug_strength}") + + if pipe_override is not None: + self.pipe = pipe_override + else: + current_device = comfy.model_management.get_torch_device() + device_str = "cuda" if current_device.type == 'cuda' else "cpu" + dtype_str = "float16" if device_str == "cuda" else "float32" + self._load_pipeline(device_str, dtype_str) + + if self.pipe is None: + raise RuntimeError("ComfyUI-NormalCrafter: Pipeline could not be loaded.") + + pil_frames = self.tensor_to_pil_list(images) + if not pil_frames: return (torch.zeros_like(images),) + + resized_pil_frames = self.resize_pil_images(pil_frames, max_res_dimension) + + # --- Temporal Padding Fix --- + num_actual_frames = len(resized_pil_frames) + effective_frames_for_pipeline = list(resized_pil_frames) # Make a mutable copy + + if num_actual_frames == 0: + print("ComfyUI-NormalCrafter: Warning - No frames to process after resizing.") + return (torch.zeros_like(images),) + + if num_actual_frames < window_size: + print(f"ComfyUI-NormalCrafter: Number of frames ({num_actual_frames}) is less than window_size ({window_size}).") + print(f"ComfyUI-NormalCrafter: Temporally padding frames by duplicating the last frame.") + padding_needed = window_size - num_actual_frames + last_frame_to_duplicate = resized_pil_frames[-1] # Use the actual last frame for padding + for _ in range(padding_needed): + effective_frames_for_pipeline.append(last_frame_to_duplicate) + print(f"ComfyUI-NormalCrafter: Frames temporally padded to {len(effective_frames_for_pipeline)} frames for pipeline processing.") + # --- End Temporal Padding Fix --- + + generator_device = self.pipe.device + generator = torch.Generator(device=generator_device).manual_seed(seed) + + print(f"ComfyUI-NormalCrafter: Processing {len(effective_frames_for_pipeline)} frames (effective) with seed {seed}. Original frames: {num_actual_frames}.") + # The log statement for the removed parameters can be updated or removed + print(f"ComfyUI-NormalCrafter: Using (internal defaults) fps={default_fps_for_time_ids}, motion_id={default_motion_bucket_id}, noise_aug={default_noise_aug_strength}") + + + pbar = comfy.utils.ProgressBar(len(effective_frames_for_pipeline)) + + with torch.inference_mode(): + output_frames_np = self.pipe( + images=effective_frames_for_pipeline, + decode_chunk_size=decode_chunk_size, + time_step_size=time_step_size, + window_size=window_size, + fps=default_fps_for_time_ids, # <<< USE DEFAULT + motion_bucket_id=default_motion_bucket_id, # <<< USE DEFAULT + noise_aug_strength=default_noise_aug_strength, # <<< USE DEFAULT + generator=generator + ).frames[0] + + # --- Truncate output if padding occurred --- + if len(effective_frames_for_pipeline) > num_actual_frames: + print(f"ComfyUI-NormalCrafter: Truncating processed frames from {output_frames_np.shape[0]} back to {num_actual_frames}.") + output_frames_np = output_frames_np[:num_actual_frames, :, :, :] + # --- End Truncation --- + + output_normals_0_1 = (output_frames_np.clip(-1., 1.) * 0.5) + 0.5 + output_tensor = torch.from_numpy(output_normals_0_1).float() + + print(f"ComfyUI-NormalCrafter: Processing complete. Output tensor shape: {output_tensor.shape}") + return (output_tensor,) \ No newline at end of file diff --git a/normalcrafter/__init__.py b/normalcrafter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/normalcrafter/normal_crafter_ppl.py b/normalcrafter/normal_crafter_ppl.py new file mode 100644 index 0000000..f8fbc17 --- /dev/null +++ b/normalcrafter/normal_crafter_ppl.py @@ -0,0 +1,578 @@ +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Union + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F +import math + +from diffusers.utils import BaseOutput, logging +from diffusers.utils.torch_utils import is_compiled_module, randn_tensor +from diffusers import DiffusionPipeline +from diffusers.pipelines.stable_video_diffusion.pipeline_stable_video_diffusion import StableVideoDiffusionPipelineOutput, StableVideoDiffusionPipeline +from PIL import Image +import cv2 + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +class NormalCrafterPipeline(StableVideoDiffusionPipeline): + + def _encode_image(self, image, device, num_videos_per_prompt, do_classifier_free_guidance, scale=1, image_size=None): + dtype = next(self.image_encoder.parameters()).dtype + + if not isinstance(image, torch.Tensor): + image = self.video_processor.pil_to_numpy(image) # (0, 255) -> (0, 1) + image = self.video_processor.numpy_to_pt(image) # (n, h, w, c) -> (n, c, h, w) + + # We normalize the image before resizing to match with the original implementation. + # Then we unnormalize it after resizing. + pixel_values = image + B, C, H, W = pixel_values.shape + patches = [pixel_values] + # patches = [] + for i in range(1, scale): + num_patches_HW_this_level = i + 1 + patch_H = H // num_patches_HW_this_level + 1 + patch_W = W // num_patches_HW_this_level + 1 + for j in range(num_patches_HW_this_level): + for k in range(num_patches_HW_this_level): + patches.append(pixel_values[:, :, j*patch_H:(j+1)*patch_H, k*patch_W:(k+1)*patch_W]) + + def encode_image(image): + image = image * 2.0 - 1.0 + if image_size is not None: + image = _resize_with_antialiasing(image, image_size) + else: + image = _resize_with_antialiasing(image, (224, 224)) + image = (image + 1.0) / 2.0 + + # Normalize the image with for CLIP input + image = self.feature_extractor( + images=image, + do_normalize=True, + do_center_crop=False, + do_resize=False, + do_rescale=False, + return_tensors="pt", + ).pixel_values + + image = image.to(device=device, dtype=dtype) + image_embeddings = self.image_encoder(image).image_embeds + if len(image_embeddings.shape) < 3: + image_embeddings = image_embeddings.unsqueeze(1) + return image_embeddings + + image_embeddings = [] + for patch in patches: + image_embeddings.append(encode_image(patch)) + image_embeddings = torch.cat(image_embeddings, dim=1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + # import pdb + # pdb.set_trace() + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_image_embeddings = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_image_embeddings, image_embeddings]) + + return image_embeddings + + def ecnode_video_vae(self, images, chunk_size: int = 14): + if isinstance(images, list): + width, height = images[0].size + else: + height, width = images[0].shape[:2] + needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast + if needs_upcasting: + self.vae.to(dtype=torch.float32) + + device = self._execution_device + images = self.video_processor.preprocess_video(images, height=height, width=width).to(device, self.vae.dtype) # torch type in range(-1, 1) with (1,3,h,w) + images = images.squeeze(0) # from (1, c, t, h, w) -> (c, t, h, w) + images = images.permute(1,0,2,3) # c, t, h, w -> (t, c, h, w) + + video_latents = [] + # chunk_size = 14 + for i in range(0, images.shape[0], chunk_size): + video_latents.append(self.vae.encode(images[i : i + chunk_size]).latent_dist.mode()) + image_latents = torch.cat(video_latents) + + # cast back to fp16 if needed + if needs_upcasting: + self.vae.to(dtype=torch.float16) + + return image_latents + + def pad_image(self, images, scale=64): + def get_pad(newW, W): + pad_W = (newW - W) // 2 + if W % 2 == 1: + pad_Ws = [pad_W, pad_W + 1] + else: + pad_Ws = [pad_W, pad_W] + return pad_Ws + + if type(images[0]) is np.ndarray: + H, W = images[0].shape[:2] + else: + W, H = images[0].size + + if W % scale == 0 and H % scale == 0: + return images, None + newW = int(np.ceil(W / scale) * scale) + newH = int(np.ceil(H / scale) * scale) + + pad_Ws = get_pad(newW, W) + pad_Hs = get_pad(newH, H) + + new_images = [] + for image in images: + if type(image) is np.ndarray: + image = cv2.copyMakeBorder(image, *pad_Hs, *pad_Ws, cv2.BORDER_CONSTANT, value=(1.,1.,1.)) + new_images.append(image) + else: + image = np.array(image) + image = cv2.copyMakeBorder(image, *pad_Hs, *pad_Ws, cv2.BORDER_CONSTANT, value=(255,255,255)) + new_images.append(Image.fromarray(image)) + return new_images, pad_Hs+pad_Ws + + def unpad_image(self, v, pad_HWs): + t, b, l, r = pad_HWs + if t > 0 or b > 0: + v = v[:, :, t:-b] + if l > 0 or r > 0: + v = v[:, :, :, l:-r] + return v + + @torch.no_grad() + def __call__( + self, + images: Union[PIL.Image.Image, List[PIL.Image.Image], torch.FloatTensor], + decode_chunk_size: Optional[int] = None, + time_step_size: Optional[int] = 1, + window_size: Optional[int] = 1, + fps: int = 7, # <<< CHANGED: Added default, will be overridden by node + motion_bucket_id: int = 127, # <<< CHANGED: Added default, will be overridden by node + noise_aug_strength: float = 0.0, # <<< CHANGED: Added default, will be overridden by node + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + return_dict: bool = True + ): + images, pad_HWs = self.pad_image(images) + + # 0. Default height and width to unet + width, height = images[0].size + num_frames = len(images) + + # 1. Check inputs. Raise error if not correct + self.check_inputs(images, height, width) + + # 2. Define call parameters + batch_size = 1 + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + self._guidance_scale = 1.0 # NormalCrafter seems to operate without CFG for normals directly + num_videos_per_prompt = 1 + do_classifier_free_guidance = False # NormalCrafter specific + num_inference_steps = 1 # For direct normal generation, this is typically 1 step if not iterative refinement + # fps, motion_bucket_id, noise_aug_strength are now passed as arguments + + output_type = "np" # Default output type from SVD pipeline is numpy array + # data_keys = ["normal"] # Not used directly in this simplified flow + use_linear_merge = True + determineTrain = True # This seems to indicate a direct generation rather than denoising from pure noise + encode_image_scale = 1 + encode_image_WH = None + + decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else 7 # Default if not provided + + # 3. Encode input image using using clip. (num_image * num_videos_per_prompt, 1, 1024) + image_embeddings = self._encode_image(images, device, num_videos_per_prompt, do_classifier_free_guidance=do_classifier_free_guidance, scale=encode_image_scale, image_size=encode_image_WH) + # 4. Encode input image using VAE + image_latents = self.ecnode_video_vae(images, chunk_size=decode_chunk_size).to(image_embeddings.dtype) + + # image_latents [num_frames, channels, height, width] ->[1, num_frames, channels, height, width] + image_latents = image_latents.unsqueeze(0) + + # 5. Get Added Time IDs + added_time_ids = self._get_add_time_ids( + fps, # <<< USES PASSED-IN VALUE + motion_bucket_id, # <<< USES PASSED-IN VALUE + noise_aug_strength, # <<< USES PASSED-IN VALUE + image_embeddings.dtype, + batch_size, + num_videos_per_prompt, + do_classifier_free_guidance, + ) + added_time_ids = added_time_ids.to(device) + + # get Start and End frame idx for each window + def get_ses(num_frames_local): # Renamed to avoid conflict with outer num_frames + ses = [] + for i in range(0, num_frames_local, time_step_size): + ses.append([i, i + window_size]) + num_to_remain = 0 + for se_idx, se_val in enumerate(ses): # Use enumerate for clarity + if se_val[1] > num_frames_local: + # Adjust the last window to not exceed num_frames_local if it's shorter than window_size + if se_idx > 0 and ses[se_idx-1][1] < num_frames_local : # Ensure there's a previous valid window + ses[se_idx] = [num_frames_local - window_size, num_frames_local] + if ses[se_idx][0] < 0: # Handle very short videos + ses[se_idx][0] = 0 + num_to_remain = se_idx +1 + else: # If this is the only window and it's too short or calculation leads to invalid range + ses[se_idx] = [0, num_frames_local] # Process what's available + num_to_remain = se_idx + 1 + + break # Stop adding more windows + num_to_remain += 1 + ses = ses[:num_to_remain] + + # Ensure the last window covers the end if not already + if ses and ses[-1][1] < num_frames_local and num_frames_local >= window_size : + new_start = num_frames_local - window_size + if new_start > ses[-1][0]: # Add only if it's a new, valid window + ses.append([new_start, num_frames_local]) + elif new_start < ses[-1][0] and new_start >= 0: # If it overlaps but starts earlier, replace last + ses[-1] = [new_start, num_frames_local] + + + # Handle case where num_frames_local < window_size + if not ses and num_frames_local > 0: + ses.append([0, num_frames_local]) # Process all available frames + + return ses + + ses = get_ses(num_frames) # Use the original num_frames from the input 'images' + + pred = None + for i, se in enumerate(ses): + # Adjust window_num_frames based on the actual slice, esp. for the last potentially shorter window + current_window_num_frames = se[1] - se[0] + window_image_embeddings = image_embeddings[se[0]:se[1]] + window_image_latents = image_latents[:, se[0]:se[1]] + # added_time_ids might need adjustment if its batch size depends on num_frames, + # but SVD usually computes it once for the whole sequence properties (fps, motion_bucket_id) + window_added_time_ids = added_time_ids # Assuming added_time_ids are constant per generation call + + if i == 0 or time_step_size >= window_size: # Corrected condition: no overlap if step >= window + to_replace_latents = None + else: + last_se = ses[i-1] + overlap_start_in_current_window = 0 + overlap_start_in_previous_pred = last_se[1] - (window_size - time_step_size) # More robust overlap calculation + num_to_replace_latents = window_size - time_step_size # This is the overlap size + + if num_to_replace_latents > 0 and pred is not None and pred.shape[1] >= (se[0] + num_to_replace_latents) : + # The slice from `pred` needs to correspond to the overlapping part of the *previous* window's output + # that aligns with the *start* of the current window's processing. + # `se[0]` is the global start index of the current window. + # `last_se[1]` is the global end index of the previous window. + # The overlap is from `se[0]` to `last_se[1]`. + # So, from `pred`, we need from `se[0]` up to `num_to_replace_latents` frames into the current window. + to_replace_latents = pred[:, se[0] : se[0] + num_to_replace_latents] + else: + to_replace_latents = None + + + latents_current_window = self.generate( + num_inference_steps, + device, + batch_size, + num_videos_per_prompt, + current_window_num_frames, # Use actual frames in this window + height, # Original height/width after padding + width, # Original height/width after padding + window_image_embeddings, + generator, + determineTrain, + to_replace_latents, # This is for conditioning the start of the current window + do_classifier_free_guidance, + window_image_latents, + window_added_time_ids + ) + + # Merge latents + if pred is None: + pred = latents_current_window + else: + if to_replace_latents is not None and use_linear_merge and time_step_size < window_size: + num_overlap_frames = window_size - time_step_size + weight = torch.linspace(1.0, 0.0, num_overlap_frames + 2, device=device, dtype=latents_current_window.dtype)[1:-1] + weight = weight.view(1, -1, 1, 1, 1) # Reshape for broadcasting + + # The part of `pred` to merge with is from global index `se[0]` for `num_overlap_frames` + pred_overlap_part = pred[:, se[0] : se[0] + num_overlap_frames] + # The part of `latents_current_window` to merge is its beginning + current_window_overlap_part = latents_current_window[:, :num_overlap_frames] + + merged_overlap = pred_overlap_part * weight + current_window_overlap_part * (1.0 - weight) + + # Update `pred` and append the new, non-overlapping part of `latents_current_window` + new_part_of_current_window = latents_current_window[:, num_overlap_frames:] + pred = torch.cat([pred[:, :se[0]], merged_overlap, new_part_of_current_window], dim=1) + else: + # No overlap or no merge, just append the new part + # This assumes pred ends exactly where the new window (minus overlap) begins + pred = torch.cat([pred[:, :se[0] + time_step_size], latents_current_window[:, window_size - time_step_size:] ], dim=1) + # A simpler, but potentially less robust way if windows aren't perfectly aligned: + # pred = torch.cat([pred[:, :se[0]], latents_current_window], dim=1) + # This needs careful handling if just appending parts to avoid duplicated or missing frames. + # For now, let's assume `pred` is built up to `se[0]` and we append `latents_current_window` + # This logic needs to be very precise for stitching. + # A safer approach if `time_step_size` is not less than `window_size`: + if time_step_size >= window_size : + pred = torch.cat([pred, latents_current_window], dim=1) # Non-overlapping windows + else: # Overlapping windows but not merging (or linear merge disabled) - take new window's data + pred = torch.cat([pred[:, :se[0]], latents_current_window], dim=1) + + + + # Ensure pred is not longer than num_frames due to windowing logic + if pred is not None and pred.shape[1] > num_frames: + pred = pred[:, :num_frames] + + + if not output_type == "latent": + needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast + if needs_upcasting: + self.vae.to(dtype=torch.float16) + + def decode_latents(latents_to_decode, num_frames_to_decode, decode_chunk_size_local): + # Ensure latents_to_decode is not None and has frames + if latents_to_decode is None or latents_to_decode.shape[1] == 0: + # Return a dummy tensor or raise an error, based on expected behavior + # For now, let's assume if latents are None, frames should be too. + # This part of SVD pipeline expects latents to be (1, T, C, H, W) + # If num_frames_to_decode is 0, it might also cause issues. + print("Warning: No latents to decode.") + # Create a dummy black frame sequence matching expected output structure + # This depends on what `self.video_processor.postprocess_video` expects. + # Assuming it handles (T, H, W, C) numpy arrays. + # Placeholder: return empty array or array of zeros. + # The original SVD decode_latents doesn't explicitly handle None input. + # Let's try to return a zero array matching expected shape if num_frames_to_decode > 0 + if num_frames_to_decode > 0: + dummy_h = latents_to_decode.shape[3] * self.vae_scale_factor if latents_to_decode is not None else 256 # example H + dummy_w = latents_to_decode.shape[4] * self.vae_scale_factor if latents_to_decode is not None else 256 # example W + return np.zeros((num_frames_to_decode, dummy_h, dummy_w, 3), dtype=np.float32) + return np.array([]) + + + frames_decoded = self.decode_latents(latents_to_decode, num_frames_to_decode, decode_chunk_size_local) + frames_decoded = self.video_processor.postprocess_video(video=frames_decoded, output_type="np") + frames_decoded = frames_decoded * 2 - 1 + return frames_decoded + + frames = decode_latents(pred, num_frames, decode_chunk_size) + if pad_HWs is not None: + frames = self.unpad_image(frames, pad_HWs) + else: + frames = pred + + self.maybe_free_model_hooks() + + if not return_dict: + return frames + + return StableVideoDiffusionPipelineOutput(frames=frames) + + + def generate( + self, + num_inference_steps, + device, + batch_size, + num_videos_per_prompt, + num_frames, + height, + width, + image_embeddings, + generator, + determineTrain, + to_replace_latents, + do_classifier_free_guidance, + image_latents, + added_time_ids, + latents=None, + ): + # 6. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 7. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_frames, + num_channels_latents, + height, + width, + image_embeddings.dtype, + device, + generator, + latents, + ) + if determineTrain: + latents[...] = 0. + + # 8. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # replace part of latents with conditons. ToDo: t embedding should also replace + if to_replace_latents is not None: + num_img_condition = to_replace_latents.shape[1] + if not determineTrain: + _noise = randn_tensor(to_replace_latents.shape, generator=generator, device=device, dtype=image_embeddings.dtype) + noisy_to_replace_latents = self.scheduler.add_noise(to_replace_latents, _noise, t.unsqueeze(0)) + latents[:, :num_img_condition] = noisy_to_replace_latents + else: + latents[:, :num_img_condition] = to_replace_latents + + + # 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 = self.scheduler.scale_model_input(latent_model_input, t) + timestep = t + # Concatenate image_latents over channels dimention + latent_model_input = torch.cat([latent_model_input, image_latents], dim=2) + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + timestep, + encoder_hidden_states=image_embeddings, + added_time_ids=added_time_ids, + return_dict=False, + )[0] + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + scheduler_output = self.scheduler.step(noise_pred, t, latents) + latents = scheduler_output.prev_sample + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + return latents +# resizing utils +# TODO: clean up later +def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True): + h, w = input.shape[-2:] + factors = (h / size[0], w / size[1]) + + # First, we have to determine sigma + # Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171 + sigmas = ( + max((factors[0] - 1.0) / 2.0, 0.001), + max((factors[1] - 1.0) / 2.0, 0.001), + ) + + # Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma + # https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206 + # But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now + ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3)) + + # Make sure it is odd + if (ks[0] % 2) == 0: + ks = ks[0] + 1, ks[1] + + if (ks[1] % 2) == 0: + ks = ks[0], ks[1] + 1 + + input = _gaussian_blur2d(input, ks, sigmas) + + output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners) + return output + + +def _compute_padding(kernel_size): + """Compute padding tuple.""" + # 4 or 6 ints: (padding_left, padding_right,padding_top,padding_bottom) + # https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad + if len(kernel_size) < 2: + raise AssertionError(kernel_size) + computed = [k - 1 for k in kernel_size] + + # for even kernels we need to do asymmetric padding :( + out_padding = 2 * len(kernel_size) * [0] + + for i in range(len(kernel_size)): + computed_tmp = computed[-(i + 1)] + + pad_front = computed_tmp // 2 + pad_rear = computed_tmp - pad_front + + out_padding[2 * i + 0] = pad_front + out_padding[2 * i + 1] = pad_rear + + return out_padding + + +def _filter2d(input, kernel): + # prepare kernel + b, c, h, w = input.shape + tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype) + + tmp_kernel = tmp_kernel.expand(-1, c, -1, -1) + + height, width = tmp_kernel.shape[-2:] + + padding_shape: list[int] = _compute_padding([height, width]) + input = torch.nn.functional.pad(input, padding_shape, mode="reflect") + + # kernel and input tensor reshape to align element-wise or batch-wise params + tmp_kernel = tmp_kernel.reshape(-1, 1, height, width) + input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1)) + + # convolve the tensor with the kernel. + output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1) + + out = output.view(b, c, h, w) + return out + + +def _gaussian(window_size: int, sigma): + if isinstance(sigma, float): + sigma = torch.tensor([[sigma]]) + + batch_size = sigma.shape[0] + + x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1) + + if window_size % 2 == 0: + x = x + 0.5 + + gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0))) + + return gauss / gauss.sum(-1, keepdim=True) + + +def _gaussian_blur2d(input, kernel_size, sigma): + if isinstance(sigma, tuple): + sigma = torch.tensor([sigma], dtype=input.dtype) + else: + sigma = sigma.to(dtype=input.dtype) + + ky, kx = int(kernel_size[0]), int(kernel_size[1]) + bs = sigma.shape[0] + kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1)) + kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1)) + out_x = _filter2d(input, kernel_x[..., None, :]) + out = _filter2d(out_x, kernel_y[..., None]) + + return out diff --git a/normalcrafter/unet.py b/normalcrafter/unet.py new file mode 100644 index 0000000..5b9123b --- /dev/null +++ b/normalcrafter/unet.py @@ -0,0 +1,368 @@ +from diffusers import UNetSpatioTemporalConditionModel +from diffusers.models.unets.unet_spatio_temporal_condition import UNetSpatioTemporalConditionOutput +from diffusers.utils import is_torch_version +import torch +from typing import Any, Dict, Optional, Tuple, Union + +def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward +CKPT_KWARGS = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + + +class DiffusersUNetSpatioTemporalConditionModelNormalCrafter(UNetSpatioTemporalConditionModel): + + @staticmethod + def forward_crossattn_down_block_dino( + module, + hidden_states: torch.Tensor, + temb: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + image_only_indicator: Optional[torch.Tensor] = None, + dino_down_block_res_samples = None, + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + output_states = () + self = module + blocks = list(zip(self.resnets, self.attentions)) + for resnet, attn in blocks: + if self.training and self.gradient_checkpointing: # TODO + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + **CKPT_KWARGS, + ) + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn), + hidden_states, + encoder_hidden_states, + image_only_indicator, + False, + **CKPT_KWARGS, + )[0] + else: + hidden_states = resnet( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + return_dict=False, + )[0] + + if dino_down_block_res_samples is not None: + hidden_states += dino_down_block_res_samples.pop(0) + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + if dino_down_block_res_samples is not None: + hidden_states += dino_down_block_res_samples.pop(0) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + @staticmethod + def forward_down_block_dino( + module, + hidden_states: torch.Tensor, + temb: Optional[torch.Tensor] = None, + image_only_indicator: Optional[torch.Tensor] = None, + dino_down_block_res_samples = None, + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + self = module + output_states = () + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + use_reentrant=False, + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + ) + else: + hidden_states = resnet( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + if dino_down_block_res_samples is not None: + hidden_states += dino_down_block_res_samples.pop(0) + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + if dino_down_block_res_samples is not None: + hidden_states += dino_down_block_res_samples.pop(0) + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + added_time_ids: torch.Tensor, + return_dict: bool = True, + image_controlnet_down_block_res_samples = None, + image_controlnet_mid_block_res_sample = None, + dino_down_block_res_samples = None, + + ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]: + r""" + The [`UNetSpatioTemporalConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`. + added_time_ids: (`torch.FloatTensor`): + The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal + embeddings and added to the time embeddings. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead of a plain + tuple. + Returns: + [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + if not hasattr(self, "custom_gradient_checkpointing"): + self.custom_gradient_checkpointing = False + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + batch_size, num_frames = sample.shape[:2] + if len(timesteps.shape) == 1: + timesteps = timesteps.expand(batch_size) + else: + timesteps = timesteps.reshape(batch_size * num_frames) + t_emb = self.time_proj(timesteps) # (B, C) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb) # (B, C) + + time_embeds = self.add_time_proj(added_time_ids.flatten()) + time_embeds = time_embeds.reshape((batch_size, -1)) + time_embeds = time_embeds.to(emb.dtype) + aug_emb = self.add_embedding(time_embeds) + if emb.shape[0] == 1: + emb = emb + aug_emb + # Repeat the embeddings num_video_frames times + # emb: [batch, channels] -> [batch * frames, channels] + emb = emb.repeat_interleave(num_frames, dim=0) + else: + aug_emb = aug_emb.repeat_interleave(num_frames, dim=0) + emb = emb + aug_emb + + # Flatten the batch and frames dimensions + # sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width] + sample = sample.flatten(0, 1) + + # encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels] + # here, our encoder_hidden_states is [batch * frames, 1, channels] + + if not sample.shape[0] == encoder_hidden_states.shape[0]: + encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0) + # 2. pre-process + sample = self.conv_in(sample) + + image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device) + + if dino_down_block_res_samples is not None: + dino_down_block_res_samples = [x for x in dino_down_block_res_samples] + sample += dino_down_block_res_samples.pop(0) + + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if dino_down_block_res_samples is None: + if self.custom_gradient_checkpointing: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + sample, res_samples = torch.utils.checkpoint.checkpoint( + create_custom_forward(downsample_block), + sample, + emb, + encoder_hidden_states, + image_only_indicator, + **CKPT_KWARGS, + ) + else: + sample, res_samples = torch.utils.checkpoint.checkpoint( + create_custom_forward(downsample_block), + sample, + emb, + image_only_indicator, + **CKPT_KWARGS, + ) + else: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + image_only_indicator=image_only_indicator, + ) + else: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + sample, res_samples = self.forward_crossattn_down_block_dino( + downsample_block, + sample, + emb, + encoder_hidden_states, + image_only_indicator, + dino_down_block_res_samples, + ) + else: + sample, res_samples = self.forward_down_block_dino( + downsample_block, + sample, + emb, + image_only_indicator, + dino_down_block_res_samples, + ) + down_block_res_samples += res_samples + + if image_controlnet_down_block_res_samples is not None: + new_down_block_res_samples = () + + for down_block_res_sample, image_controlnet_down_block_res_sample in zip( + down_block_res_samples, image_controlnet_down_block_res_samples + ): + down_block_res_sample = (down_block_res_sample + image_controlnet_down_block_res_sample) / 2 + new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) + + down_block_res_samples = new_down_block_res_samples + + # 4. mid + if self.custom_gradient_checkpointing: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), + sample, + emb, + encoder_hidden_states, + image_only_indicator, + **CKPT_KWARGS, + ) + else: + sample = self.mid_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + + if image_controlnet_mid_block_res_sample is not None: + sample = (sample + image_controlnet_mid_block_res_sample) / 2 + + # 5. up + mid_up_block_out_samples = [sample, ] + down_block_out_sampels = [] + for i, upsample_block in enumerate(self.up_blocks): + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + down_block_out_sampels.append(res_samples[-1]) + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + if self.custom_gradient_checkpointing: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(upsample_block), + sample, + res_samples, + emb, + encoder_hidden_states, + image_only_indicator, + **CKPT_KWARGS + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + else: + if self.custom_gradient_checkpointing: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(upsample_block), + sample, + res_samples, + emb, + image_only_indicator, + **CKPT_KWARGS + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + image_only_indicator=image_only_indicator, + ) + mid_up_block_out_samples.append(sample) + # 6. post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + if self.custom_gradient_checkpointing: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.conv_out), + sample, + **CKPT_KWARGS + ) + else: + sample = self.conv_out(sample) + + # 7. Reshape back to original shape + sample = sample.reshape(batch_size, num_frames, *sample.shape[1:]) + + if not return_dict: + return (sample, down_block_out_sampels[::-1], mid_up_block_out_samples) + + return UNetSpatioTemporalConditionOutput(sample=sample) \ No newline at end of file diff --git a/normalcrafter/utils.py b/normalcrafter/utils.py new file mode 100644 index 0000000..b848c7e --- /dev/null +++ b/normalcrafter/utils.py @@ -0,0 +1,64 @@ +from typing import Union, List +import tempfile +import numpy as np +import PIL.Image +import matplotlib.cm as cm +import mediapy +import torch +from decord import VideoReader, cpu + + +def read_video_frames(video_path, process_length, target_fps, max_res): + print("==> processing video: ", video_path) + vid = VideoReader(video_path, ctx=cpu(0)) + print("==> original video shape: ", (len(vid), *vid.get_batch([0]).shape[1:])) + original_height, original_width = vid.get_batch([0]).shape[1:3] + + if max(original_height, original_width) > max_res: + scale = max_res / max(original_height, original_width) + height = round(original_height * scale) + width = round(original_width * scale) + else: + height = original_height + width = original_width + + vid = VideoReader(video_path, ctx=cpu(0), width=width, height=height) + + fps = vid.get_avg_fps() if target_fps == -1 else target_fps + stride = round(vid.get_avg_fps() / fps) + stride = max(stride, 1) + frames_idx = list(range(0, len(vid), stride)) + print( + f"==> downsampled shape: {len(frames_idx), *vid.get_batch([0]).shape[1:]}, with stride: {stride}" + ) + if process_length != -1 and process_length < len(frames_idx): + frames_idx = frames_idx[:process_length] + print( + f"==> final processing shape: {len(frames_idx), *vid.get_batch([0]).shape[1:]}" + ) + frames = vid.get_batch(frames_idx).asnumpy().astype(np.uint8) + frames = [PIL.Image.fromarray(x) for x in frames] + + return frames, fps + +def save_video( + video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], + output_video_path: str = None, + fps: int = 10, + crf: int = 18, +) -> str: + if output_video_path is None: + output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name + + if isinstance(video_frames[0], np.ndarray): + video_frames = [(frame * 255).astype(np.uint8) for frame in video_frames] + + elif isinstance(video_frames[0], PIL.Image.Image): + video_frames = [np.array(frame) for frame in video_frames] + mediapy.write_video(output_video_path, video_frames, fps=fps, crf=crf) + return output_video_path + +def vis_sequence_normal(normals: np.ndarray): + normals = normals.clip(-1., 1.) + normals = normals * 0.5 + 0.5 + return normals diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..abf2abb --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +mediapy +decord + + +# diffusers +# If you do not have diffusers installed, you need to install it. +# You can typically install a compatible version by running the following command inside your activated VENV: +# pip install diffusers +# Repeat this process for any other packages you may be missing that show as errors during bootup or runtime \ No newline at end of file