726 lines
24 KiB
Python
726 lines
24 KiB
Python
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,
|
|
}
|