Merge pull request #13 from kaski23/main
Better Performance in Pre-and Postprocessing steps
This commit is contained in:
AIWarper
+78
-111
@@ -2,6 +2,7 @@
|
||||
|
||||
import torch
|
||||
import torchvision.transforms as transforms
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import os
|
||||
@@ -10,6 +11,7 @@ import logging
|
||||
import comfy.model_management
|
||||
import model_management
|
||||
import comfy.utils
|
||||
from comfy.utils import common_upscale
|
||||
import folder_paths # ComfyUI's way to get model paths
|
||||
|
||||
# Handle PIL version compatibility
|
||||
@@ -294,60 +296,34 @@ class NormalCrafterNode:
|
||||
print(f"ComfyUI-NormalCrafter: Pipeline loaded to {target_device_for_this_run} with dtype {final_torch_dtype_for_load}.")
|
||||
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 _validate_dimensions_for_attention(self, width: int, height: int) -> tuple:
|
||||
|
||||
def make_safe_size(self, images: torch.Tensor, max_res_dim: int, device: torch.device) -> torch.Tensor:
|
||||
"""
|
||||
Validates and adjusts dimensions to be compatible with flash attention requirements.
|
||||
Returns (safe_width, safe_height, needs_adjustment) tuple.
|
||||
Skaliert ein Batch von Bildern mit Torch, ohne PIL/NumPy,
|
||||
und sorgt gleich für Attention-kompatible Dimensionen.
|
||||
"""
|
||||
# Flash attention typically requires dimensions divisible by 8 or 16
|
||||
# and the total elements to be within certain bounds
|
||||
|
||||
# After VAE encoding, dimensions will be divided by 8
|
||||
latent_width = width // 8
|
||||
latent_height = height // 8
|
||||
|
||||
# Common constraints for flash attention
|
||||
min_dim = 2 # Minimum dimension after encoding
|
||||
alignment = 8 # Dimension alignment requirement
|
||||
|
||||
# Check if dimensions are safe
|
||||
if latent_width < min_dim or latent_height < min_dim:
|
||||
# Dimensions too small
|
||||
safe_width = max(width, min_dim * 8)
|
||||
safe_height = max(height, min_dim * 8)
|
||||
return safe_width, safe_height, True
|
||||
|
||||
# Check alignment
|
||||
if latent_width % alignment != 0 or latent_height % alignment != 0:
|
||||
# Need alignment adjustment
|
||||
safe_width = ((latent_width + alignment - 1) // alignment) * alignment * 8
|
||||
safe_height = ((latent_height + alignment - 1) // alignment) * alignment * 8
|
||||
return safe_width, safe_height, True
|
||||
|
||||
return width, height, False
|
||||
imgs = images.to(device, non_blocking=True).permute(0, 3, 1, 2) # (B,C,H,W)
|
||||
_, _, H, W = imgs.shape
|
||||
|
||||
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), LANCZOS))
|
||||
return resized_images
|
||||
# 1. Scale down to max_res_dim (if necessary
|
||||
if max(H, W) > max_res_dim:
|
||||
scale = max_res_dim / max(H, W)
|
||||
target_h, target_w = int(round(H * scale)), int(round(W * scale))
|
||||
else:
|
||||
target_h, target_w = H, W
|
||||
|
||||
# 2. Enforce Attention-Safe-Dimension (minimum 64)
|
||||
min_dim = 64
|
||||
align = 64
|
||||
target_h = max(min_dim, ((target_h + align - 1) // align) * align)
|
||||
target_w = max(min_dim, ((target_w + align - 1) // align) * align)
|
||||
|
||||
# 3. Torch-Interpolation
|
||||
imgs_resized = F.interpolate(imgs, size=(target_h, target_w), mode="bilinear", align_corners=False)
|
||||
|
||||
return imgs_resized.permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
|
||||
def process(self, images: torch.Tensor, seed: int, max_res_dimension: int,
|
||||
window_size: int, time_step_size: int, decode_chunk_size: int,
|
||||
offload_pipe_to_cpu_on_finish: bool, # <<< NEW PARAMETER
|
||||
@@ -361,6 +337,7 @@ class NormalCrafterNode:
|
||||
if pipe_override is not None:
|
||||
self.pipe = pipe_override
|
||||
print("ComfyUI-NormalCrafter: Using provided pipe_override.")
|
||||
|
||||
else:
|
||||
current_comfy_device = comfy.model_management.get_torch_device()
|
||||
device_str = "cuda" if current_comfy_device.type == 'cuda' else "cpu"
|
||||
@@ -370,7 +347,7 @@ class NormalCrafterNode:
|
||||
|
||||
if self.pipe is None:
|
||||
raise RuntimeError("ComfyUI-NormalCrafter: Pipeline could not be loaded.")
|
||||
|
||||
|
||||
# Ensure the pipe instance self.pipe is on the correct device for processing *before* using it
|
||||
# This is important if self.pipe came from the global cache and might have been on CPU
|
||||
processing_device = comfy.model_management.get_torch_device()
|
||||
@@ -378,79 +355,54 @@ class NormalCrafterNode:
|
||||
print(f"ComfyUI-NormalCrafter: Moving self.pipe from {self.pipe.device} to {processing_device} for processing.")
|
||||
self.pipe.to(processing_device)
|
||||
|
||||
|
||||
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)
|
||||
#Resizing
|
||||
orig_h, orig_w = images.shape[1:3]
|
||||
resized_tensor = self.make_safe_size(images, max_res_dimension, processing_device)
|
||||
|
||||
num_actual_frames = len(resized_pil_frames)
|
||||
effective_frames_for_pipeline = list(resized_pil_frames)
|
||||
orig_h, orig_w = images.shape[1:3]
|
||||
resized_tensor = self.make_safe_size(images, max_res_dimension, processing_device)
|
||||
|
||||
# Get number of frames in the tensor
|
||||
num_actual_frames = resized_tensor.shape[0]
|
||||
|
||||
|
||||
|
||||
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}). Padding...")
|
||||
padding_needed = window_size - num_actual_frames
|
||||
last_frame_to_duplicate = resized_pil_frames[-1]
|
||||
for _ in range(padding_needed):
|
||||
effective_frames_for_pipeline.append(last_frame_to_duplicate)
|
||||
last_frame = resized_tensor[-1:] # Shape (1,H,W,C)
|
||||
pad_frames = last_frame.repeat(padding_needed, 1, 1, 1)
|
||||
effective_frames_for_pipeline = torch.cat([resized_tensor, pad_frames], dim=0)
|
||||
else:
|
||||
effective_frames_for_pipeline = resized_tensor
|
||||
|
||||
|
||||
|
||||
generator_device = self.pipe.device # Should be processing_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: {num_actual_frames}.")
|
||||
effective_frame_count = effective_frames_for_pipeline.shape[0]
|
||||
|
||||
print(f"ComfyUI-NormalCrafter: Processing {effective_frame_count} frames (effective) with seed {seed}. Original: {num_actual_frames}.")
|
||||
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)) # This pbar seems not used by the pipe.
|
||||
pbar = comfy.utils.ProgressBar(effective_frame_count) # optional, Pipe nutzt eh ihre eigene
|
||||
|
||||
# Check dimensions before processing
|
||||
if resized_pil_frames:
|
||||
first_frame_width, first_frame_height = resized_pil_frames[0].size
|
||||
safe_width, safe_height, needs_adjustment = self._validate_dimensions_for_attention(first_frame_width, first_frame_height)
|
||||
|
||||
if needs_adjustment and use_xformers != "force":
|
||||
print(f"ComfyUI-NormalCrafter: Detected potentially problematic dimensions ({first_frame_width}x{first_frame_height}). "
|
||||
f"Adjusting to safe dimensions ({safe_width}x{safe_height}) for flash attention compatibility.")
|
||||
# Re-resize ALL frames in effective_frames_for_pipeline to safe dimensions
|
||||
# This ensures padded frames also get resized
|
||||
for i in range(len(effective_frames_for_pipeline)):
|
||||
if isinstance(effective_frames_for_pipeline[i], Image.Image):
|
||||
effective_frames_for_pipeline[i] = effective_frames_for_pipeline[i].resize((safe_width, safe_height), LANCZOS)
|
||||
else:
|
||||
# Handle case where frame might not be a PIL Image
|
||||
print(f"ComfyUI-NormalCrafter: Warning - frame {i} is not a PIL Image, skipping resize")
|
||||
|
||||
# Final safety check: ensure all frames have the same dimensions
|
||||
if effective_frames_for_pipeline:
|
||||
reference_size = None
|
||||
all_same_size = True
|
||||
|
||||
for i, frame in enumerate(effective_frames_for_pipeline):
|
||||
if isinstance(frame, Image.Image):
|
||||
if reference_size is None:
|
||||
reference_size = frame.size
|
||||
elif frame.size != reference_size:
|
||||
all_same_size = False
|
||||
print(f"ComfyUI-NormalCrafter: Frame {i} has size {frame.size}, expected {reference_size}")
|
||||
|
||||
if not all_same_size and reference_size:
|
||||
print(f"ComfyUI-NormalCrafter: Ensuring all frames are {reference_size} for consistency")
|
||||
for i in range(len(effective_frames_for_pipeline)):
|
||||
if isinstance(effective_frames_for_pipeline[i], Image.Image):
|
||||
if effective_frames_for_pipeline[i].size != reference_size:
|
||||
effective_frames_for_pipeline[i] = effective_frames_for_pipeline[i].resize(reference_size, LANCZOS)
|
||||
|
||||
output_frames_np = None
|
||||
output_frames_pt = None
|
||||
error_occurred = False
|
||||
|
||||
# Try processing with current settings
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
output_frames_np = self.pipe( # self.pipe should be on processing_device here
|
||||
images=effective_frames_for_pipeline,
|
||||
output_frames_pt = self.pipe( # self.pipe should be on processing_device here
|
||||
images=resized_tensor,
|
||||
decode_chunk_size=decode_chunk_size,
|
||||
time_step_size=time_step_size,
|
||||
window_size=window_size,
|
||||
@@ -472,8 +424,8 @@ class NormalCrafterNode:
|
||||
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
output_frames_np = self.pipe(
|
||||
images=effective_frames_for_pipeline,
|
||||
output_frames_pt = self.pipe(
|
||||
images=resized_tensor,
|
||||
decode_chunk_size=decode_chunk_size,
|
||||
time_step_size=time_step_size,
|
||||
window_size=window_size,
|
||||
@@ -498,14 +450,29 @@ class NormalCrafterNode:
|
||||
else:
|
||||
raise e
|
||||
|
||||
if output_frames_np is None:
|
||||
if output_frames_pt is None:
|
||||
raise RuntimeError("ComfyUI-NormalCrafter: Failed to generate output frames.")
|
||||
|
||||
if len(effective_frames_for_pipeline) > num_actual_frames:
|
||||
output_frames_np = output_frames_np[:num_actual_frames, :, :, :]
|
||||
|
||||
output_normals_0_1 = (output_frames_np.clip(-1., 1.) * 0.5) + 0.5
|
||||
output_tensor = torch.from_numpy(output_normals_0_1).float()
|
||||
if len(resized_tensor) > num_actual_frames:
|
||||
output_frames_pt = output_frames_pt[:num_actual_frames, ...]
|
||||
|
||||
|
||||
# Output from NormalCrafter has the pattern: [B,C,H,W]
|
||||
# 1. Normalize to a range of [0,1]
|
||||
output_frames_pt = (output_frames_pt.clamp(-1., 1.) * 0.5) + 0.5
|
||||
|
||||
# 2. Resize back to input-resolution
|
||||
output_frames_pt = F.interpolate(
|
||||
output_frames_pt,
|
||||
size=(orig_h, orig_w), # Eingangsauflösung
|
||||
mode="bilinear",
|
||||
align_corners=False
|
||||
)
|
||||
|
||||
# 3. Permute Tensor to [B,H,W,C] for ComfyUI-Compatibility
|
||||
output_frames_pt = output_frames_pt.permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
|
||||
|
||||
# --- Explicit Offload After Processing ---
|
||||
# Only offload the globally managed pipe (NORMALCRAFTER_PIPE), not an overridden one.
|
||||
@@ -527,5 +494,5 @@ class NormalCrafterNode:
|
||||
except Exception as e:
|
||||
print(f"ComfyUI-NormalCrafter: Error offloading pipeline to CPU: {e}")
|
||||
|
||||
print(f"ComfyUI-NormalCrafter: Processing complete. Output tensor shape: {output_tensor.shape}")
|
||||
return (output_tensor,)
|
||||
print(f"ComfyUI-NormalCrafter: Processing complete. Output tensor shape: {output_frames_pt.shape}")
|
||||
return (output_frames_pt,)
|
||||
|
||||
@@ -18,144 +18,99 @@ 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):
|
||||
def _encode_image(self, image: torch.Tensor, 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)
|
||||
if isinstance(image, torch.Tensor):
|
||||
# Expected: (T,H,W,C) or (T,C,H,W) in [0,1]
|
||||
if image.dim() != 4:
|
||||
raise ValueError("Tensor image must be 4D (T,H,W,C) or (T,C,H,W).")
|
||||
if image.shape[-1] == 3: # (T,H,W,C) -> (T,C,H,W)
|
||||
image = image.permute(0,3,1,2)
|
||||
elif image.shape[1] == 3: # (T,C,H,W)
|
||||
pass
|
||||
else:
|
||||
raise ValueError("Channel dim must be 3.")
|
||||
|
||||
# We normalize the image before resizing to match with the original implementation.
|
||||
# Then we unnormalize it after resizing.
|
||||
pixel_values = image
|
||||
pixel_values = image.to(device=device, dtype=torch.float32) # B=T
|
||||
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])
|
||||
num_patches = i + 1
|
||||
ph = H // num_patches + 1
|
||||
pw = W // num_patches + 1
|
||||
for j in range(num_patches):
|
||||
for k in range(num_patches):
|
||||
patches.append(pixel_values[:, :, j*ph:(j+1)*ph, k*pw:(k+1)*pw])
|
||||
|
||||
def encode_image(img_bchw):
|
||||
img_bchw = img_bchw * 2.0 - 1.0
|
||||
img_bchw = _resize_with_antialiasing(img_bchw, image_size or (224, 224))
|
||||
img_bchw = (img_bchw + 1.0) / 2.0 # zurück in [0,1]
|
||||
|
||||
# CLIP-Preproc
|
||||
feat = self.feature_extractor(
|
||||
images=img_bchw, do_normalize=True, do_center_crop=False,
|
||||
do_resize=False, do_rescale=False, return_tensors="pt",
|
||||
).pixel_values.to(device=device, dtype=dtype)
|
||||
|
||||
emb = self.image_encoder(feat).image_embeds
|
||||
if emb.dim() < 3:
|
||||
emb = emb.unsqueeze(1)
|
||||
return emb
|
||||
|
||||
image_embeddings = torch.cat([encode_image(p) for p in patches], dim=1)
|
||||
else:
|
||||
raise ValueError("_encode_image() didn't get handed a torch-tensor")
|
||||
|
||||
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()
|
||||
# Duplicate
|
||||
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])
|
||||
negative = torch.zeros_like(image_embeddings)
|
||||
image_embeddings = torch.cat([negative, 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]
|
||||
|
||||
def encode_video_vae(self, images: torch.Tensor, chunk_size: int = 14):
|
||||
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)
|
||||
|
||||
if isinstance(images, torch.Tensor):
|
||||
# (T,H,W,C) oder (T,C,H,W) in [0,1] -> (1, C, T, H, W) in [-1,1]
|
||||
if images.shape[-1] == 3: # (T,H,W,C) -> (T,C,H,W)
|
||||
imgs = images.permute(0,3,1,2)
|
||||
else: # (T,C,H,W)
|
||||
imgs = images
|
||||
imgs = imgs.to(device=device, dtype=self.vae.dtype)
|
||||
imgs = imgs * 2.0 - 1.0
|
||||
imgs = imgs.permute(1,0,2,3).unsqueeze(0) # (C,T,H,W) -> (1,C,T,H,W)
|
||||
else:
|
||||
raise ValueError("encode_video_vae() didn't get handed a torch-tensor")
|
||||
|
||||
# (1,C,T,H,W) -> (C,T,H,W) -> (T,C,H,W)
|
||||
imgs = imgs.squeeze(0).permute(1,0,2,3)
|
||||
|
||||
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)
|
||||
for i in range(0, imgs.shape[0], chunk_size):
|
||||
video_latents.append(self.vae.encode(imgs[i:i+chunk_size]).latent_dist.mode())
|
||||
image_latents = torch.cat(video_latents, dim=0)
|
||||
|
||||
# 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],
|
||||
images: torch.FloatTensor,
|
||||
decode_chunk_size: Optional[int] = None,
|
||||
time_step_size: Optional[int] = 1,
|
||||
window_size: Optional[int] = 1,
|
||||
@@ -165,14 +120,28 @@ class NormalCrafterPipeline(StableVideoDiffusionPipeline):
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
return_dict: bool = True
|
||||
):
|
||||
images, pad_HWs = self.pad_image(images)
|
||||
#images, pad_HWs = self.pad_image(images)
|
||||
|
||||
# 0. Default height and width to unet
|
||||
width, height = images[0].size
|
||||
num_frames = len(images)
|
||||
if isinstance(images, torch.Tensor):
|
||||
if images.dim() != 4:
|
||||
raise ValueError("images tensor must be 4D (T,H,W,C) or (T,C,H,W)")
|
||||
if images.shape[-1] == 3: # (T,H,W,C)
|
||||
T, H, W, C = images.shape
|
||||
images_for_enc = images.permute(0,3,1,2) # (T,C,H,W)
|
||||
elif images.shape[1] == 3: # (T,C,H,W)
|
||||
T, C, H, W = images.shape[0], images.shape[1], images.shape[2], images.shape[3]
|
||||
images_for_enc = images
|
||||
else:
|
||||
raise ValueError("Channel dim must be 3.")
|
||||
num_frames, height, width = T, H, W
|
||||
else:
|
||||
width, height = images[0].size
|
||||
num_frames = len(images)
|
||||
images_for_enc = images # bleibt Liste von PIL
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(images, height, width)
|
||||
# 1. Check inputs. Raise error if not correct -> Redundant, due to the rescaling done in normal_crafter_node.py
|
||||
# self.check_inputs(images, height, width)
|
||||
|
||||
# 2. Define call parameters
|
||||
batch_size = 1
|
||||
@@ -186,7 +155,7 @@ class NormalCrafterPipeline(StableVideoDiffusionPipeline):
|
||||
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
|
||||
output_type = "pt" # Default output type from SVD pipeline is numpy array, but we use torch because we are cooler
|
||||
# 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
|
||||
@@ -198,7 +167,7 @@ class NormalCrafterPipeline(StableVideoDiffusionPipeline):
|
||||
# 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 = self.encode_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)
|
||||
@@ -369,13 +338,12 @@ class NormalCrafterPipeline(StableVideoDiffusionPipeline):
|
||||
|
||||
|
||||
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 = self.video_processor.postprocess_video(video=frames_decoded, output_type=output_type) # "pt"
|
||||
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user