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

Added CPU Offloading

Browse Source
This commit is contained in:
AIWarper
2025-05-09 15:52:50 -06:00
parent 6b7f64cbde
commit d749003ebb
2 changed files with 73 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2 -2
View File
@@ -14,7 +14,7 @@ It allows you to generate temporally consistent normal map sequences from input
2. Install the required dependencies. Navigate to the `ComfyUI-NormalCrafter` directory and install using the `requirements.txt`:
```bash
ACTIVATE YOUR VENV FIRST
cd ComfyUI-NormalCrafter
ComfyUI-NormalCrafterWrapper
pip install -r requirements.txt
```
(See the `requirements.txt` for notes on the `diffusers` dependency, which ComfyUI often manages.)
@@ -54,4 +54,4 @@ You may be able to increase `window_size` and `time_step_size` substantially (e.
* `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.
Refer to `requirements.txt` for more details.
normal_crafter_nodes.py
+71 -45
View File
@@ -27,7 +27,7 @@ except ImportError as e:
# Global variable to cache the pipeline
NORMALCRAFTER_PIPE = None
CURRENT_PIPE_CONFIG = {}
CURRENT_PIPE_CONFIG = {} # Stores the config {"device": "cuda"/"cpu", "dtype": "float16"/"float32"}
# Define model paths and repo ID
NORMALCRAFTER_REPO_ID = "Yanrui95/NormalCrafter"
@@ -49,9 +49,7 @@ class NormalCrafterNode:
"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
"offload_pipe_to_cpu_on_finish": ("BOOLEAN", {"default": True}), # <<< NEW INPUT
},
"optional": {
"pipe_override": ("NORMALCRAFTER_PIPE",),
@@ -93,18 +91,28 @@ class NormalCrafterNode:
def _load_pipeline(self, device_str="cuda", dtype_str="float16"):
global NORMALCRAFTER_PIPE, CURRENT_PIPE_CONFIG
# This is the configuration requested for the *current* run
requested_config = {
"device": device_str,
"dtype": dtype_str,
}
# If pipe exists and its last known config matches what's requested for *this* run
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)
target_device_for_this_run = comfy.model_management.get_torch_device() if device_str == "cuda" else torch.device("cpu")
# Safety check: ensure it's actually on the target_device_for_this_run.
# This handles cases where ComfyUI might have moved it, or if our offload didn't update CURRENT_PIPE_CONFIG.
# However, with proper CURRENT_PIPE_CONFIG updates, this `if` might rarely trigger if CURRENT_PIPE_CONFIG already reflects "cpu".
if NORMALCRAFTER_PIPE.device != target_device_for_this_run:
print(f"ComfyUI-NormalCrafter: Moving existing pipeline from {NORMALCRAFTER_PIPE.device} to {target_device_for_this_run}.")
NORMALCRAFTER_PIPE.to(target_device_for_this_run)
self.pipe = NORMALCRAFTER_PIPE
return self.pipe
# If pipe doesn't exist, or its last known config (CURRENT_PIPE_CONFIG)
# is different from what's requested for this run (e.g., it was on CPU, now GPU is requested),
# then we proceed to load/reload.
print("ComfyUI-NormalCrafter: Loading NormalCrafter pipeline...")
local_nc_model_path = self._download_normalcrafter_model_if_needed()
@@ -114,11 +122,11 @@ class NormalCrafterNode:
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
local_nc_model_path, subfolder="unet", low_cpu_mem_usage=True, torch_dtype=torch_dtype # Initially load to CPU
)
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
local_nc_model_path, subfolder="vae", low_cpu_mem_usage=True, torch_dtype=torch_dtype # Initially load to CPU
)
svd_xt_variant = "fp16" if torch_dtype == torch.float16 else None
@@ -127,18 +135,19 @@ class NormalCrafterNode:
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")
target_device_for_this_run = 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.")
# Consider pipe.enable_attention_slicing() as a fallback if VRAM is an issue during processing
pipe.to(target_device)
pipe.to(target_device_for_this_run) # Move the newly loaded pipe to the target device
NORMALCRAFTER_PIPE = pipe
CURRENT_PIPE_CONFIG = requested_config
CURRENT_PIPE_CONFIG = requested_config # Update global config to reflect the state of the loaded pipe
self.pipe = NORMALCRAFTER_PIPE
print(f"ComfyUI-NormalCrafter: Pipeline loaded to {target_device} with dtype {torch_dtype}.")
print(f"ComfyUI-NormalCrafter: Pipeline loaded to {target_device_for_this_run} with dtype {torch_dtype}.")
return self.pipe
def tensor_to_pil_list(self, images_tensor: torch.Tensor) -> list:
@@ -165,82 +174,99 @@ class NormalCrafterNode:
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
offload_pipe_to_cpu_on_finish: bool, # <<< NEW PARAMETER
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
print("ComfyUI-NormalCrafter: Using provided 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"
current_comfy_device = comfy.model_management.get_torch_device()
device_str = "cuda" if current_comfy_device.type == 'cuda' else "cpu"
# Determine dtype based on device (float16 for CUDA, float32 for CPU)
dtype_str = "float16" if device_str == "cuda" and comfy.model_management.should_use_fp16() else "float32"
self._load_pipeline(device_str, dtype_str)
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()
if self.pipe.device != processing_device:
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)
# --- Temporal Padding Fix ---
num_actual_frames = len(resized_pil_frames)
effective_frames_for_pipeline = list(resized_pil_frames) # Make a mutable copy
effective_frames_for_pipeline = list(resized_pil_frames)
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.")
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] # Use the actual last frame for padding
last_frame_to_duplicate = resized_pil_frames[-1]
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_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 frames: {num_actual_frames}.")
# The log statement for the removed parameters can be updated or removed
print(f"ComfyUI-NormalCrafter: Processing {len(effective_frames_for_pipeline)} 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))
pbar = comfy.utils.ProgressBar(len(effective_frames_for_pipeline)) # This pbar seems not used by the pipe.
with torch.inference_mode():
output_frames_np = self.pipe(
output_frames_np = self.pipe( # self.pipe should be on processing_device here
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
fps=default_fps_for_time_ids,
motion_bucket_id=default_motion_bucket_id,
noise_aug_strength=default_noise_aug_strength,
generator=generator
# SVD pipeline has its own progress bar, no need to pass pbar here
).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()
# --- Explicit Offload After Processing ---
# Only offload the globally managed pipe (NORMALCRAFTER_PIPE), not an overridden one.
# And only if no pipe_override was used for this run.
if pipe_override is None and \
offload_pipe_to_cpu_on_finish and \
NORMALCRAFTER_PIPE is not None and \
NORMALCRAFTER_PIPE.device.type == 'cuda': # Check if it's on CUDA before moving
print("ComfyUI-NormalCrafter: Offloading globally cached pipeline to CPU after processing.")
try:
NORMALCRAFTER_PIPE.to("cpu")
global CURRENT_PIPE_CONFIG # Make sure to get the global
if "device" in CURRENT_PIPE_CONFIG:
CURRENT_PIPE_CONFIG["device"] = "cpu" # Update its state to "cpu"
else: # Should not happen if config is always set
CURRENT_PIPE_CONFIG = {"device": "cpu", "dtype": CURRENT_PIPE_CONFIG.get("dtype", "float32")}
comfy.model_management.soft_empty_cache() # Ask ComfyUI to try and free VRAM
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,)
return (output_tensor,)