AIWarper-ComfyUI-NormalCrafterWrapper/normal_crafter_nodes.py

# ComfyUI/custom_nodes/ComfyUI-NormalCrafter/normal_crafter_nodes.py

import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
import numpy as np
from PIL import Image
import os
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
try:
    LANCZOS = Image.Resampling.LANCZOS
except AttributeError:
    LANCZOS = Image.LANCZOS

# 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 = {} # Stores the config {"device": "cuda"/"cpu", "dtype": "float16"/"float32"}

# 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 DetailTransfer:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "target": ("IMAGE", ),
                "source": ("IMAGE", ),
                "mode": ([
                    "add",
                    "multiply",
                    "screen",
                    "overlay",
                    "soft_light",
                    "hard_light",
                    "color_dodge",
                    "color_burn",
                    "difference",
                    "exclusion",
                    "divide",
                    
                    ], 
                    {"default": "add"}
                    ),
                "blur_sigma": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.01}),
                "blend_factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001,  "round": 0.001}),
            },
            "optional": {
                "mask": ("MASK", ),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process"
    CATEGORY = "NormalCrafter"

    def adjust_mask(self, mask, target_tensor):
        # Add a channel dimension and repeat to match the channel number of the target tensor
        if len(mask.shape) == 3:
            mask = mask.unsqueeze(1)  # Add a channel dimension
            target_channels = target_tensor.shape[1]
            mask = mask.expand(-1, target_channels, -1, -1)  # Expand the channel dimension to match the target tensor's channels
    
        return mask


    def process(self, target, source, mode, blur_sigma, blend_factor, mask=None):
        B, H, W, C = target.shape
        device = model_management.get_torch_device()
        target_tensor = target.permute(0, 3, 1, 2).clone().to(device)
        source_tensor = source.permute(0, 3, 1, 2).clone().to(device)

        if target.shape[1:] != source.shape[1:]:
            source_tensor = comfy.utils.common_upscale(source_tensor, W, H, "bilinear", "disabled")

        if source.shape[0] < B:
            source = source[0].unsqueeze(0).repeat(B, 1, 1, 1)

        kernel_size = int(6 * int(blur_sigma) + 1)

        gaussian_blur = transforms.GaussianBlur(kernel_size=(kernel_size, kernel_size), sigma=(blur_sigma, blur_sigma))

        blurred_target = gaussian_blur(target_tensor)
        blurred_source = gaussian_blur(source_tensor)
        
        if mode == "add":
            tensor_out = (source_tensor - blurred_source) + blurred_target
        elif mode == "multiply":
            tensor_out = source_tensor * blurred_target
        elif mode == "screen":
            tensor_out = 1 - (1 - source_tensor) * (1 - blurred_target)
        elif mode == "overlay":
            tensor_out = torch.where(blurred_target < 0.5, 2 * source_tensor * blurred_target, 1 - 2 * (1 - source_tensor) * (1 - blurred_target))
        elif mode == "soft_light":
            tensor_out = (1 - 2 * blurred_target) * source_tensor**2 + 2 * blurred_target * source_tensor
        elif mode == "hard_light":
            tensor_out = torch.where(source_tensor < 0.5, 2 * source_tensor * blurred_target, 1 - 2 * (1 - source_tensor) * (1 - blurred_target))
        elif mode == "difference":
            tensor_out = torch.abs(blurred_target - source_tensor)
        elif mode == "exclusion":
            tensor_out = 0.5 - 2 * (blurred_target - 0.5) * (source_tensor - 0.5)
        elif mode == "color_dodge":
            tensor_out = blurred_target / (1 - source_tensor)
        elif mode == "color_burn":
            tensor_out = 1 - (1 - blurred_target) / source_tensor
        elif mode == "divide":
            tensor_out = (source_tensor / blurred_source) * blurred_target
        else:
            tensor_out = source_tensor
        
        tensor_out = torch.lerp(target_tensor, tensor_out, blend_factor)
        if mask is not None:
            # Call the function and pass in mask and target_tensor
            mask = self.adjust_mask(mask, target_tensor)
            mask = mask.to(device)
            tensor_out = torch.lerp(target_tensor, tensor_out, mask)
        tensor_out = torch.clamp(tensor_out, 0, 1)
        tensor_out = tensor_out.permute(0, 2, 3, 1).cpu().float()
        return (tensor_out,)

class NormalCrafterNode:
    def __init__(self):
        self.pipe = None # Instance variable for the pipeline
        self.last_processed_dimensions = None # Track last processed dimensions

    @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}),
                "offload_pipe_to_cpu_on_finish": ("BOOLEAN", {"default": True}), # <<< NEW INPUT
                "use_xformers": (["auto", "disable", "force"], {"default": "auto"}), # <<< NEW INPUT
            },
            "optional": {
                "pipe_override": ("NORMALCRAFTER_PIPE",),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process"
    CATEGORY = "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 _configure_xformers(self, pipe, use_xformers="auto"):
        """
        Configures xformers memory efficient attention for a given pipeline.
        Handles 'auto', 'disable', and 'force' modes.
        """
        if use_xformers == "disable":
            if hasattr(pipe, 'disable_xformers_memory_efficient_attention'):
                print("ComfyUI-NormalCrafter: Xformers memory efficient attention disabled.")
                try:
                    pipe.disable_xformers_memory_efficient_attention()
                except Exception as e:
                    print(f"ComfyUI-NormalCrafter: Warning - could not disable xformers: {e}")
            else:
                print("ComfyUI-NormalCrafter: Pipeline does not support disabling xformers.")
        elif use_xformers == "force":
            if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
                print("ComfyUI-NormalCrafter: Xformers memory efficient attention forced on.")
                pipe.enable_xformers_memory_efficient_attention()
            else:
                print("ComfyUI-NormalCrafter: Warning - Pipeline does not support xformers, but 'force' was requested.")
        else: # "auto" or any other value
            if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
                try:
                    pipe.enable_xformers_memory_efficient_attention()
                    print("ComfyUI-NormalCrafter: Xformers memory efficient attention enabled (auto).")
                except Exception as e:
                    print(f"ComfyUI-NormalCrafter: Failed to enable Xformers memory efficient attention: {e}. Using standard attention.")
                    # If xformers fails, we can still use standard attention
                    pass
            else:
                print("ComfyUI-NormalCrafter: Pipeline does not support xformers. Using standard attention.")

    def _load_pipeline(self, device_str_requested="cuda", dtype_str_requested="float16", use_xformers="auto"):
        global NORMALCRAFTER_PIPE, CURRENT_PIPE_CONFIG

        target_device_for_this_run = comfy.model_management.get_torch_device() if device_str_requested == "cuda" else torch.device("cpu")
        
        # Determine the torch_dtype we want for this run
        if dtype_str_requested == "float16": final_torch_dtype_for_load = torch.float16
        elif dtype_str_requested == "bf16": final_torch_dtype_for_load = torch.bfloat16
        else: final_torch_dtype_for_load = torch.float32

        if NORMALCRAFTER_PIPE is not None:
            # Pipe exists. Check if its current dtype and target device are suitable.
            pipe_actual_dtype = NORMALCRAFTER_PIPE.dtype # The true dtype of the existing pipe
            pipe_actual_dtype_str = "float16" if pipe_actual_dtype == torch.float16 else \
                                    "bf16" if pipe_actual_dtype == torch.bfloat16 else "float32"

            if dtype_str_requested == pipe_actual_dtype_str:
                # Dtypes match. Just check device.
                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) 
                    # The warning about fp16 on cpu will appear here if target_device_for_this_run is cpu, 
                    # but we're moving an fp16 pipe from cpu to gpu in the typical reload case, which is fine.
                
                CURRENT_PIPE_CONFIG = {"device": str(target_device_for_this_run), "dtype": dtype_str_requested}
                self.pipe = NORMALCRAFTER_PIPE
                # Configure xformers for existing pipe
                self._configure_xformers(NORMALCRAFTER_PIPE, use_xformers)
                # print(f"ComfyUI-NormalCrafter: Reusing existing pipeline. Now on {target_device_for_this_run} with {dtype_str_requested}.")
                return self.pipe
            else:
                # Dtype mismatch (e.g., user wants float32 now, but pipe is float16). Must reload fully.
                print(f"ComfyUI-NormalCrafter: Requested dtype {dtype_str_requested} differs from existing pipe's dtype {pipe_actual_dtype_str}. Reloading pipeline.")
                NORMALCRAFTER_PIPE = None # Force a full reload by clearing the global pipe
        
        # If NORMALCRAFTER_PIPE is None (either initially, or forced by dtype mismatch above)
        print("ComfyUI-NormalCrafter: Loading/Re-loading NormalCrafter pipeline from scratch...")
        local_nc_model_path = self._download_normalcrafter_model_if_needed()

        # Load components with final_torch_dtype_for_load
        unet = DiffusersUNetSpatioTemporalConditionModelNormalCrafter.from_pretrained(
            local_nc_model_path, subfolder="unet", low_cpu_mem_usage=True, torch_dtype=final_torch_dtype_for_load
        )
        vae = AutoencoderKLTemporalDecoder.from_pretrained(
            local_nc_model_path, subfolder="vae", low_cpu_mem_usage=True, torch_dtype=final_torch_dtype_for_load
        )
        svd_xt_variant = "fp16" if final_torch_dtype_for_load == torch.float16 else None

        pipe = NormalCrafterPipeline.from_pretrained(
            SVD_XT_REPO_ID, unet=unet, vae=vae, torch_dtype=final_torch_dtype_for_load, variant=svd_xt_variant,
        )

        # Configure xformers with our robust handler
        self._configure_xformers(pipe, use_xformers)

        pipe.to(target_device_for_this_run)
        NORMALCRAFTER_PIPE = pipe
        CURRENT_PIPE_CONFIG = {
            "device": str(target_device_for_this_run),
            "dtype": dtype_str_requested # The dtype it's configured with for this run
        }
        self.pipe = NORMALCRAFTER_PIPE
        print(f"ComfyUI-NormalCrafter: Pipeline loaded to {target_device_for_this_run} with dtype {final_torch_dtype_for_load}.")
        return self.pipe

    
    def make_safe_size(self, images: torch.Tensor, max_res_dim: int, device: torch.device) -> torch.Tensor:
        """
        Skaliert ein Batch von Bildern mit Torch, ohne PIL/NumPy,
        und sorgt gleich für Attention-kompatible Dimensionen.
        """
        imgs = images.to(device, non_blocking=True).permute(0, 3, 1, 2)  # (B,C,H,W)
        _, _, H, W = imgs.shape

        # 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
                use_xformers: str = "auto", # <<< NEW PARAMETER
                pipe_override=None):

        default_fps_for_time_ids = 7
        default_motion_bucket_id = 127
        default_noise_aug_strength = 0.0

        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"
            # 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, use_xformers)

        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)

        #Resizing
        orig_h, orig_w = images.shape[1:3]
        resized_tensor = self.make_safe_size(images, max_res_dimension, processing_device)
        
        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 = 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)

        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(effective_frame_count)  # optional, Pipe nutzt eh ihre eigene



        output_frames_pt = None
        error_occurred = False
        
        # Try processing with current settings
        try:
            with torch.inference_mode():
                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,
                    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]
        except RuntimeError as e:
            if "CUDA error" in str(e) and "invalid configuration argument" in str(e) and use_xformers != "disable":
                print(f"ComfyUI-NormalCrafter: Flash attention error encountered: {str(e)}")
                error_occurred = True
                
                # Try disabling xformers and retrying
                if hasattr(self.pipe, 'disable_xformers_memory_efficient_attention'):
                    print("ComfyUI-NormalCrafter: Attempting to disable xformers and retry...")
                    self.pipe.disable_xformers_memory_efficient_attention()
                    
                    try:
                        with torch.inference_mode():
                            output_frames_pt = self.pipe(
                                images=resized_tensor,
                                decode_chunk_size=decode_chunk_size,
                                time_step_size=time_step_size,
                                window_size=window_size,
                                fps=default_fps_for_time_ids,
                                motion_bucket_id=default_motion_bucket_id,
                                noise_aug_strength=default_noise_aug_strength,
                                generator=generator
                            ).frames[0]
                        print("ComfyUI-NormalCrafter: Successfully processed without xformers.")
                    except Exception as retry_error:
                        print(f"ComfyUI-NormalCrafter: Retry without xformers also failed: {retry_error}")
                        raise retry_error
                    
                    # Re-enable xformers for future runs if it was auto mode
                    if use_xformers == "auto":
                        try:
                            self.pipe.enable_xformers_memory_efficient_attention()
                        except:
                            pass
                else:
                    raise e
            else:
                raise e
        
        if output_frames_pt is None:
            raise RuntimeError("ComfyUI-NormalCrafter: Failed to generate output frames.")

        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.
        # 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_frames_pt.shape}")
        return (output_frames_pt,)