nodes/AIWarper-ComfyUI-NormalCrafterWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit of ComfyUI-NormalCrafter node

Browse Source
This commit is contained in:
AIWarper
2025-05-09 12:18:59 -06:00
commit f2ece13a51
10 changed files with 1384 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+42
View File
@@ -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
README.md
+56
View File
@@ -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.
__init__.py
+21
View File
@@ -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 = {}
git
View File
normal_crafter_nodes.py
+246
View File
@@ -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,)
normalcrafter/__init__.py
View File
normalcrafter/normal_crafter_ppl.py
+578
View File
@@ -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
normalcrafter/unet.py
+368
View File
@@ -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)
normalcrafter/utils.py
+64
View File
@@ -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
requirements.txt
+9
View File
@@ -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