import os import sys sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import torch import logging import numpy as np import cv2 from PIL import Image from huggingface_hub import snapshot_download import torchvision.transforms as T from LightGlue.lightglue import LightGlue, SuperPoint from LightGlue.lightglue.utils import rbd from models_diffusers.unet_spatio_temporal_condition import ( UNetSpatioTemporalConditionModel, ) from models_diffusers.controlnet_svd import ControlNetSVDModel try: from cotracker.predictor import CoTrackerPredictor, sample_trajectories_with_ref except: pass from lineart_extractor.canny import CannyDetector from lineart_extractor.hed import HEDdetector from lineart_extractor.lineart import LineartDetector from lineart_extractor.lineart_anime import LineartAnimeDetector from pipelines.AniDoc import AniDocPipeline from comfy.utils import ProgressBar from comfy.model_management import XFORMERS_IS_AVAILABLE import folder_paths log = logging.getLogger("AniDoc") DIFFUSERS_DIR = os.path.join(folder_paths.models_dir, "diffusers") ANIDOC_DIR = os.path.join(DIFFUSERS_DIR, "anidoc") ANIDOC_CONTROLNET_DIR = os.path.join(ANIDOC_DIR, "controlnet") SVD_I2V_DIR = os.path.join( DIFFUSERS_DIR, "stable-video-diffusion-img2vid-xt", ) COTRACKER_DIR = os.path.join(folder_paths.models_dir, "cotracker") if "cotracker" not in folder_paths.folder_names_and_paths: current_paths = [COTRACKER_DIR] else: current_paths, _ = folder_paths.folder_names_and_paths["cotracker"] folder_paths.folder_names_and_paths["cotracker"] = ( current_paths, folder_paths.supported_pt_extensions, ) class AniDocLoader: @classmethod def INPUT_TYPES(cls): return { "required": {}, } RETURN_TYPES = ("ANIDOC_PIPELINE",) RETURN_NAMES = ("anidoc_pipeline",) OUTPUT_TOOLTIPS = ("AniDoc Pipeline",) CATEGORY = "AniDoc" FUNCTION = "load_anidoc" DESCRIPTION = "Load AniDoc pipeline" def __init__(self): self.pipeline = None def load_anidoc( self, anidoc_path=ANIDOC_DIR, svd_img2vid_path=SVD_I2V_DIR, controlnet_path=ANIDOC_CONTROLNET_DIR, device="cuda", dtype=torch.float16, ): if self.pipeline is None: pbar = ProgressBar(5) pbar.update(1) log.info(f"Loading model from: {anidoc_path}") log.info("Missing models will be downloaded") try: snapshot_download( repo_id="Yhmeng1106/anidoc", ignore_patterns=["*.md"], local_dir=DIFFUSERS_DIR, local_dir_use_symlinks=False, ) except: log.info("Couldn't download models") unet = UNetSpatioTemporalConditionModel.from_pretrained( anidoc_path, subfolder="unet", torch_dtype=dtype, low_cpu_mem_usage=True, custom_resume=True, ) unet.to(device, dtype) pbar.update(1) log.info(f"Loading controlnet from: {controlnet_path}") controlnet = ControlNetSVDModel.from_pretrained(controlnet_path) controlnet.to(device, dtype) pbar.update(1) if XFORMERS_IS_AVAILABLE: log.info("Enabling XFormers") unet.enable_xformers_memory_efficient_attention() log.info(f"Loading model from: {svd_img2vid_path}") log.info("Missing models will be downloaded") try: snapshot_download( repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1", allow_patterns=["*.json", "*fp16*"], ignore_patterns=["*unet*"], local_dir=svd_img2vid_path, local_dir_use_symlinks=False, ) except: log.info("Couldn't download models") pbar.update(1) pipeline = AniDocPipeline.from_pretrained( svd_img2vid_path, unet=unet, controlnet=controlnet, low_cpu_mem_usage=False, torch_dtype=dtype, variant="fp16", ) pipeline.to(device) self.pipeline = pipeline pbar.update(1) return (self.pipeline,) class LoadAniDocCoTracker: @classmethod def INPUT_TYPES(cls): return { "required": { "tracking": ("BOOLEAN", {"default": True}), "cotracker_model": (folder_paths.get_filename_list("cotracker"),), "tracker_shift_grid": ([0, 1], {"default": 0}), "tracker_grid_size": ( "INT", {"default": 8, "min": 2, "max": 32, "step": 2}, ), "tracker_grid_query_frame": ("INT", {"default": 0}), "tracker_backward_tracking": ("BOOLEAN", {"default": False}), "max_points": ( "INT", {"default": 50, "min": 10, "max": 100, "step": 5}, ), }, } RETURN_TYPES = ("ANIDOC_COTRACKER",) RETURN_NAMES = ("cotracker",) OUTPUT_TOOLTIPS = ("CoTracker",) CATEGORY = "AniDoc" FUNCTION = "load_tracker" DESCRIPTION = "Load CoTracker for AniDoc" def __init__(self): self.tracker = None self.tracker_shift_grid = None def load_tracker( self, cotracker_model, tracking, tracker_shift_grid, tracker_grid_size, tracker_grid_query_frame, tracker_backward_tracking, max_points, device="cuda", dtype=torch.float32, ): try: import cotracker except: raise ImportError("Couldn't import cotracker module. Please install it to use this node") if tracking: if self.tracker is None or self.tracker_shift_grid != tracker_shift_grid: cotracker_model_path = folder_paths.get_full_path( "cotracker", cotracker_model ) log.info(f"Loading tracker model from {cotracker_model_path}") tracker = CoTrackerPredictor( checkpoint=cotracker_model_path, shift_grid=tracker_shift_grid, ) tracker.requires_grad_(False) tracker.to(device, dtype=dtype) self.tracker = tracker self.tracker_shift_grid = tracker_shift_grid return ( { "tracking": tracking, "tracker": self.tracker, "grid_size": self.tracker_shift_grid, "grid_query_frame": tracker_grid_query_frame, "backward_tracking": tracker_backward_tracking, "max_points": max_points, }, ) else: return ( { "tracking": False, "tracker": None, "grid_size": tracker_grid_size, "grid_query_frame": tracker_grid_query_frame, "backward_tracking": tracker_backward_tracking, "max_points": max_points, }, ) class GetAniDocControlnetImages: @classmethod def INPUT_TYPES(cls): return { "required": { "input_images": ("IMAGE",), "lineart_detector": ( ["none", "canny", "hed", "lineart", "lineart_anime"], {"default": "lineart"}, ), "sketch_quantization": ("BOOLEAN", {"default": True}), "width": ("INT", {"default": 512, "min": 64, "max": 1024, "step": 8}), "height": ("INT", {"default": 320, "min": 64, "max": 1024, "step": 8}), "device": (["cpu", "cuda"], {"default": "cuda"}), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("controlnet_images",) OUTPUT_TOOLTIPS = ("Processed controlnet images",) CATEGORY = "AniDoc" FUNCTION = "get_controlnet_images" DESCRIPTION = "Get lineart controlnet images for AniDoc" def load_lineart_detector(self, lineart_detector, device): if lineart_detector == "none": return None if lineart_detector == "canny": return CannyDetector() elif lineart_detector == "hed": return HEDdetector() elif lineart_detector == "lineart": return LineartDetector(device) elif lineart_detector == "lineart_anime": return LineartAnimeDetector(device) def invert_images(self, image): if np.unique(image).size == 2: return cv2.bitwise_not(image) else: return 255 - image def get_controlnet_images( self, input_images, lineart_detector, sketch_quantization, width, height, device="cuda", ): pbar = ProgressBar(len(input_images) + 1) log.info(f"Loading lineart detector: {lineart_detector}") detector = self.load_lineart_detector(lineart_detector, device) pbar.update(1) log.info("Processing images with lineart detector") controlnet_images = [] for img_tensor in input_images: sketch = (img_tensor * 255.0).clamp(0, 255).byte().cpu().numpy() if detector is not None: if lineart_detector == "canny": sketch = detector(sketch, 100, 200) else: sketch = detector(sketch) if lineart_detector in ["canny", "hed"]: sketch = self.invert_images(sketch) if len(sketch.shape) == 2: sketch = np.repeat(sketch[:, :, np.newaxis], 3, axis=2) if not sketch_quantization: sketch = (sketch > 200).astype(np.uint8) * 255 sketch = torch.nn.functional.interpolate( torch.from_numpy(sketch).permute(2, 0, 1).unsqueeze(0).float() / 255.0, size=(height, width), mode="bilinear", align_corners=False, ).squeeze(0) controlnet_images.append(sketch) pbar.update(1) controlnet_images = torch.stack(controlnet_images) controlnet_images = controlnet_images.permute(0, 2, 3, 1) return (controlnet_images,) class AniDocSampler: @classmethod def INPUT_TYPES(cls): return { "required": { "anidoc_pipeline": ("ANIDOC_PIPELINE",), "controlnet_images": ("IMAGE",), "reference_image": ("IMAGE",), "repeat_matching": ("BOOLEAN", {"default": False}), "fps": ("INT", {"default": 7, "min": 7, "max": 100, "step": 1}), "steps": ("INT", {"default": 25, "min": 1, "max": 10000, "step": 1}), "noise_aug": ( "FLOAT", {"default": 0.02, "min": 0.0, "max": 10.0, "step": 0.01}, ), "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), "motion_bucket_id": ( "INT", {"default": 127, "min": 1, "max": 300, "step": 1}, ), "decode_chunk_size": ( "INT", {"default": 8, "min": 1, "max": 256, "step": 1}, ), }, "optional": { "cotracker": ("ANIDOC_COTRACKER",), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("video_frames",) OUTPUT_TOOLTIPS = ("Video Frames",) CATEGORY = "AniDoc" FUNCTION = "sample" DESCRIPTION = "Sampler for AniDoc" def safe_round(self, coords, size): height, width = size[1], size[2] rounded_coords = np.round(coords).astype(int) rounded_coords[:, 0] = np.clip(rounded_coords[:, 0], 0, width - 1) rounded_coords[:, 1] = np.clip(rounded_coords[:, 1], 0, height - 1) return rounded_coords def generate_point_map(self, size, coords0, coords1): h, w = size[1], size[2] mask0 = np.zeros((h, w), dtype=np.uint8) mask1 = np.zeros((h, w), dtype=np.uint8) for i, (coord0, coord1) in enumerate(zip(coords0, coords1)): x0, y0 = int(round(coord0[0])), int(round(coord0[1])) x1, y1 = int(round(coord1[0])), int(round(coord1[1])) if 0 <= x0 < w and 0 <= y0 < h: mask0[y0, x0] = i + 1 if 0 <= x1 < w and 0 <= y1 < h: mask1[y1, x1] = i + 1 return mask0, mask1 def select_multiple_points(self, points0, points1, num_points): N = len(points0) num_points = min(num_points, N) indices = np.random.choice(N, size=num_points, replace=False) return points0[indices], points1[indices] def generate_point_map_frames(self, size, coords0, coords1, visibility): h, w = size[1], size[2] mask0 = np.zeros((h, w), dtype=np.uint8) num_frames = coords1.shape[0] mask1 = np.zeros((num_frames, h, w), dtype=np.uint8) for i, coord0 in enumerate(coords0): x0, y0 = int(round(coord0[0])), int(round(coord0[1])) if 0 <= x0 < w and 0 <= y0 < h: mask0[y0, x0] = i + 1 for frame_idx in range(num_frames): coords_frame = coords1[frame_idx] for i, coord1 in enumerate(coords_frame): x1, y1 = int(round(coord1[0])), int(round(coord1[1])) if 0 <= x1 < w and 0 <= y1 < h and visibility[frame_idx, i]: mask1[frame_idx, y1, x1] = i + 1 return mask0, mask1 def get_conditioning( self, tracker, extractor, matcher, controlnet_images, reference_image, frames, width, height, repeat_matching, tracking, tracker_grid_size, tracker_grid_query_frame, tracker_backward_tracking, max_points, device, dtype=torch.float16, ): controlnet_sketch_condition = [ T.ToTensor()(img).unsqueeze(0) for img in controlnet_images ] controlnet_sketch_condition = ( torch.cat(controlnet_sketch_condition, dim=0) .unsqueeze(0) .to(device, dtype=dtype) ) controlnet_sketch_condition = (controlnet_sketch_condition - 0.5) / 0.5 with torch.no_grad(): ref_img_value = ( T.ToTensor()(reference_image).to(device, dtype=dtype).to(torch.float32) ) current_img = ( T.ToTensor()(controlnet_images[0]) .to(device, dtype=dtype) .to(torch.float32) ) feats0 = extractor.extract(ref_img_value) feats1 = extractor.extract(current_img) matches01 = matcher({"image0": feats0, "image1": feats1}) feats0, feats1, matches01 = [rbd(x) for x in [feats0, feats1, matches01]] matches = matches01["matches"] points0 = feats0["keypoints"][matches[..., 0]] points1 = feats1["keypoints"][matches[..., 1]] points0 = points0.cpu().numpy() points1 = points1.cpu().numpy() points0 = self.safe_round(points0, current_img.shape) points1 = self.safe_round(points1, current_img.shape) num_points = min(50, points0.shape[0]) points0, points1 = self.select_multiple_points(points0, points1, num_points) mask1, mask2 = self.generate_point_map( size=current_img.shape, coords0=points0, coords1=points1 ) point_map1 = torch.from_numpy(mask1) point_map2 = torch.from_numpy(mask2) point_map1 = ( point_map1.unsqueeze(0) .unsqueeze(0) .unsqueeze(0) .to(device, dtype=dtype) ) point_map2 = ( point_map2.unsqueeze(0) .unsqueeze(0) .unsqueeze(0) .to(device, dtype=dtype) ) point_map = torch.cat([point_map1, point_map2], dim=2) conditional_pixel_values = ref_img_value.unsqueeze(0).unsqueeze(0) conditional_pixel_values = (conditional_pixel_values - 0.5) / 0.5 point_map_with_ref = torch.cat([point_map, conditional_pixel_values], dim=2) original_shape = list(point_map_with_ref.shape) new_shape = original_shape.copy() new_shape[1] = frames - 1 if repeat_matching: matching_controlnet_image = point_map_with_ref.repeat( 1, frames, 1, 1, 1 ) controlnet_condition = torch.cat( [controlnet_sketch_condition, matching_controlnet_image], dim=2 ) elif tracking: with torch.no_grad(): video_for_tracker = ( controlnet_sketch_condition * 0.5 + 0.5 ) * 255.0 queries = np.insert(points1, 0, 0, axis=1) queries = ( torch.from_numpy(queries).to(device, torch.float).unsqueeze(0) ) if queries.shape[1] == 0: mask1 = np.zeros((height, width), dtype=np.uint8) mask2 = np.zeros((frames, height, width), dtype=np.uint8) else: pred_tracks, pred_visibility = tracker( video_for_tracker.to(dtype=torch.float32), queries=queries, grid_size=tracker_grid_size, grid_query_frame=tracker_grid_query_frame, backward_tracking=tracker_backward_tracking, ) ( pred_tracks_sampled, pred_visibility_sampled, points0_sampled, ) = sample_trajectories_with_ref( pred_tracks.cpu(), pred_visibility.cpu(), torch.from_numpy(points0).unsqueeze(0).cpu(), max_points=max_points, motion_threshold=1, vis_threshold=3, ) if pred_tracks_sampled is None: mask1 = np.zeros((height, width), dtype=np.uint8) mask2 = np.zeros((frames, height, width), dtype=np.uint8) else: pred_tracks_sampled = ( pred_tracks_sampled.squeeze(0).cpu().numpy() ) pred_visibility_sampled = ( pred_visibility_sampled.squeeze(0).cpu().numpy() ) points0_sampled = points0_sampled.squeeze(0).cpu().numpy() for frame_id in range(frames): pred_tracks_sampled[frame_id] = self.safe_round( pred_tracks_sampled[frame_id], current_img.shape ) points0_sampled = self.safe_round( points0_sampled, current_img.shape ) mask1, mask2 = self.generate_point_map_frames( size=current_img.shape, coords0=points0_sampled, coords1=pred_tracks_sampled, visibility=pred_visibility_sampled, ) point_map1 = torch.from_numpy(mask1) point_map2 = torch.from_numpy(mask2) point_map1 = ( point_map1.unsqueeze(0) .unsqueeze(0) .repeat(1, frames, 1, 1, 1) .to(device, dtype=dtype) ) point_map2 = ( point_map2.unsqueeze(0).unsqueeze(2).to(device, dtype=dtype) ) point_map = torch.cat([point_map1, point_map2], dim=2) conditional_pixel_values_repeat = conditional_pixel_values.repeat( 1, frames, 1, 1, 1 ) point_map_with_ref = torch.cat( [point_map, conditional_pixel_values_repeat], dim=2 ) controlnet_condition = torch.cat( [controlnet_sketch_condition, point_map_with_ref], dim=2 ) else: zero_tensor = torch.zeros(new_shape).to(device, dtype=dtype) matching_controlnet_image = torch.cat( (point_map_with_ref, zero_tensor), dim=1 ) controlnet_condition = torch.cat( [controlnet_sketch_condition, matching_controlnet_image], dim=2 ) return controlnet_condition def sample( self, anidoc_pipeline, controlnet_images, reference_image, repeat_matching=False, cotracker={ "tracking": False, "tracker": None, "grid_size": 8, "grid_query_frame": 0, "backward_tracking": False, "max_points": 50, }, fps=7, steps=25, noise_aug=0.02, seed=0, motion_bucket_id=127, decode_chunk_size=8, device="cuda", dtype=torch.float16, ): extractor = SuperPoint(max_num_keypoints=2000).eval().to(device) matcher = LightGlue(features="superpoint").eval().to(device) width, height = controlnet_images.shape[2], controlnet_images.shape[1] reference_image = reference_image.permute(0, 3, 1, 2) reference_image = torch.nn.functional.interpolate( reference_image, size=(height, width), mode="bilinear", align_corners=False ) reference_image = reference_image.squeeze(0) reference_image = reference_image.permute(1, 2, 0).cpu().numpy() reference_image = (reference_image * 255).clip(0, 255).astype(np.uint8) reference_image = Image.fromarray(reference_image) images_np = (controlnet_images.cpu().numpy() * 255).astype(np.uint8) controlnet_images = [Image.fromarray(image) for image in images_np] pbar = ProgressBar(steps + 1) log.info("Getting controlnet conditioning for AniDoc") controlnet_condition = self.get_conditioning( cotracker["tracker"], extractor, matcher, controlnet_images, reference_image, len(controlnet_images), width, height, repeat_matching, cotracker["tracking"], cotracker["grid_size"], cotracker["grid_query_frame"], cotracker["backward_tracking"], cotracker["max_points"], device, dtype, ) pbar.update(1) generator = torch.manual_seed(seed) log.info("Generating video frames with AniDoc") with torch.inference_mode(): video_frames = anidoc_pipeline( reference_image, controlnet_condition, width=width, height=height, num_frames=len(controlnet_images), num_inference_steps=steps, motion_bucket_id=motion_bucket_id, fps=fps, noise_aug_strength=noise_aug, decode_chunk_size=decode_chunk_size, generator=generator, callback_on_step_end=lambda *args, **kwargs: (pbar.update(1), kwargs)[ -1 ], ).frames[0] tensor_frames = [T.ToTensor()(img) for img in video_frames] tensor_frames = torch.stack(tensor_frames) tensor_frames = tensor_frames.permute(0, 2, 3, 1) return (tensor_frames,) NODE_CLASS_MAPPINGS = { "AniDocLoader": AniDocLoader, "LoadCoTracker": LoadAniDocCoTracker, "GetAniDocControlnetImages": GetAniDocControlnetImages, "AniDocSampler": AniDocSampler, }