import sys
from pyparsing import col
sys.path.insert(0,".")
import argparse
from packaging import version
import glob
import os
from LightGlue.lightglue import LightGlue, SuperPoint, DISK, SIFT, ALIKED, DoGHardNet
from LightGlue.lightglue.utils import load_image, rbd
from cotracker.predictor import CoTrackerPredictor, sample_trajectories, generate_gassian_heatmap, sample_trajectories_with_ref
import torch
from diffusers.utils.import_utils import is_xformers_available
from models_diffusers.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel
from pipelines.AniDoc import AniDocPipeline
from models_diffusers.controlnet_svd import ControlNetSVDModel
from diffusers.utils import load_image, export_to_video, export_to_gif
import time
from lineart_extractor.annotator.lineart import LineartDetector
import numpy as np
from PIL import Image
from utils import load_images_from_folder,export_gif_with_ref,export_gif_side_by_side,extract_frames_from_video,safe_round,select_multiple_points,generate_point_map,generate_point_map_frames,export_gif_side_by_side_complete,export_gif_side_by_side_complete_ablation
import random
import torchvision.transforms as T
from LightGlue.lightglue import viz2d
import matplotlib.pyplot as plt
from cotracker.utils.visualizer import Visualizer, read_video_from_path
from torchvision.transforms import PILToTensor
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--pretrained_model_name_or_path", type=str, default="pretrained_weights/stable-video-diffusion-img2vid-xt", help="Path to the input image.")
parser.add_argument(
"--pretrained_unet", type=str, help="Path to the input image.",
default="pretrained_weights/anidoc"
)
parser.add_argument(
"--controlnet_model_name_or_path", type=str, help="Path to the input image.",
default="pretrained_weights/anidoc/controlnet"
)
parser.add_argument("--output_dir", type=str, default=None, help="Path to the output video.")
parser.add_argument("--seed", type=int, default=42, help="random seed.")
parser.add_argument("--noise_aug", type=float, default=0.02)
parser.add_argument("--num_frames", type=int, default=14)
parser.add_argument("--width", type=int, default=512)
parser.add_argument("--height", type=int, default=320)
parser.add_argument("--all_sketch",action="store_true",help="all_sketch")
parser.add_argument("--not_quant_sketch",action="store_true",help="not_quant_sketch")
parser.add_argument("--repeat_sketch",action="store_true",help="not_quant_sketch")
parser.add_argument("--matching",action="store_true",help="add keypoint matching")
parser.add_argument("--tracking",action="store_true",help="tracking keypoint")
parser.add_argument("--repeat_matching",action="store_true",help="not tracking, but just simply repeat")
parser.add_argument("--tracker_point_init", type=str, default='gaussion', choices=['dift', 'gaussion', 'both'], help="Regular grid size")
parser.add_argument(
"--tracker_shift_grid",
type=int, default=0, choices=[0, 1],
help="shift the grid for the tracker")
parser.add_argument("--tracker_grid_size", type=int, default=8, help="Regular grid size")
parser.add_argument(
"--tracker_grid_query_frame",
type=int,
default=0,
help="Compute dense and grid tracks starting from this frame",
)
parser.add_argument(
"--tracker_backward_tracking",
action="store_true",
help="Compute tracks in both directions, not only forward",
)
parser.add_argument("--control_image", type=str, default=None, help="Path to the output video.")
parser.add_argument("--ref_image", type=str, default=None, help="Path to the output video.")
parser.add_argument("--max_points", type=int, default=10)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = get_args()
dtype = torch.float16
unet = UNetSpatioTemporalConditionModel.from_pretrained(
args.pretrained_unet,
subfolder="unet",
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
custom_resume=True,
)
unet.to("cuda",dtype)
if args.controlnet_model_name_or_path:
controlnet = ControlNetSVDModel.from_pretrained(
args.controlnet_model_name_or_path,
)
else:
controlnet = ControlNetSVDModel.from_unet(
unet,
conditioning_channels=8
)
controlnet.to("cuda",dtype)
if is_xformers_available():
import xformers
xformers_version = version.parse(xformers.__version__)
unet.enable_xformers_memory_efficient_attention()
else:
raise ValueError(
"xformers is not available. Make sure it is installed correctly")
pipe = AniDocPipeline.from_pretrained(
args.pretrained_model_name_or_path,
unet=unet,
controlnet=controlnet,
low_cpu_mem_usage=False,
torch_dtype=torch.float16, variant="fp16"
)
pipe.to("cuda")
device = "cuda"
detector = LineartDetector(device)
extractor = SuperPoint(max_num_keypoints=2000).eval().to(device) # load the extractor
matcher = LightGlue(features='superpoint').eval().to(device) # load the matcher
tracker = CoTrackerPredictor(
checkpoint="pretrained_weights/cotracker2.pth",
shift_grid=args.tracker_shift_grid,
)
tracker.requires_grad_(False)
tracker.to(device, dtype=torch.float32)
width, height = args.width, args.height
# image = load_image('dalle3_cat.jpg')
if args.output_dir is None:
args.output_dir = "results"
os.makedirs(args.output_dir, exist_ok=True)
image_folder_list=[
'data_test/sample1.mp4',
]
ref_image_list=[
"data_test/sample1.png",
]
if args.ref_image is not None and args.control_image is not None:
ref_image_list=[args.ref_image]
image_folder_list=[args.control_image]
for val_id ,each_sample in enumerate(image_folder_list):
if os.path.isdir(each_sample):
control_images=load_images_from_folder(each_sample)
elif each_sample.endswith(".mp4"):
control_images = extract_frames_from_video(each_sample)
ref_image=load_image(ref_image_list[val_id]).resize((width, height))
#resize:
for j, each in enumerate(control_images):
control_images[j]=control_images[j].resize((width, height))
# load image from folder
if args.all_sketch:
controlnet_image=[]
for k in range(len(control_images)):
sketch=control_images[k]
sketch = np.array(sketch)
sketch=detector(sketch,coarse=False)
sketch=np.repeat(sketch[:, :, np.newaxis], 3, axis=2)
if args.not_quant_sketch:
pass
else:
sketch= (sketch > 200).astype(np.uint8)*255
sketch = Image.fromarray(sketch).resize((width, height))
controlnet_image.append(sketch)
controlnet_sketch_condition = [T.ToTensor()(img).unsqueeze(0) for img in controlnet_image]
controlnet_sketch_condition = torch.cat(controlnet_sketch_condition, dim=0).unsqueeze(0).to(device, dtype=torch.float16)
controlnet_sketch_condition = (controlnet_sketch_condition - 0.5) / 0.5 #(1,14,3,h,w)
# matching condition
with torch.no_grad():
ref_img_value = T.ToTensor()(ref_image).to(device, dtype=torch.float16) #(0,1)
ref_img_value = ref_img_value.to(torch.float32)
current_img= T.ToTensor()(controlnet_image[0]).to(device, dtype=torch.float16) #(0,1)
current_img = current_img.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()
# points0_org=points0.copy()
points1 = points1.cpu().numpy()
points0 = safe_round(points0, current_img.shape)
points1 = safe_round(points1, current_img.shape)
num_points = min(50, points0.shape[0])
points0,points1 = select_multiple_points(points0, points1, num_points)
mask1, mask2 = generate_point_map(size=current_img.shape, coords0=points0, coords1=points1)
# import ipdb;ipdb.set_trace()
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=torch.float16)
point_map2 = point_map2.unsqueeze(0).unsqueeze(0).unsqueeze(0).to(device, dtype=torch.float16)
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] = args.num_frames-1
if args.repeat_matching:
matching_controlnet_image=point_map_with_ref.repeat(1,args.num_frames,1,1,1)
controlnet_condition=torch.cat([controlnet_sketch_condition, matching_controlnet_image], dim=2)
elif args.tracking:
with torch.no_grad():
video_for_tracker = (controlnet_sketch_condition * 0.5 + 0.5) * 255.
queries = np.insert(points1,0,0,axis=1)
queries =torch.from_numpy(queries).to(device,torch.float).unsqueeze(0)
if queries.shape[1]==0:
pred_tracks_sampled=None
points0_sampled = None
else:
pred_tracks, pred_visibility = tracker(
video_for_tracker.to(dtype=torch.float32),
queries=queries,
grid_size=args.tracker_grid_size, # 8
grid_query_frame=args.tracker_grid_query_frame, # 0
backward_tracking=args.tracker_backward_tracking, # False
# segm_mask=segm_mask,
)
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=args.max_points,
motion_threshold=1,
vis_threshold=3,
)
if pred_tracks_sampled is None:
mask1 = np.zeros((args.height, args.width), dtype=np.uint8)
mask2 = np.zeros((args.num_frames,args.height, args.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(args.num_frames):
pred_tracks_sampled[frame_id] = safe_round(pred_tracks_sampled[frame_id],current_img.shape)
points0_sampled = safe_round(points0_sampled,current_img.shape)
mask1, mask2 = 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,args.num_frames,1,1,1).to(device, dtype=torch.float16)
point_map2 = point_map2.unsqueeze(0).unsqueeze(2).to(device, dtype=torch.float16)
point_map=torch.cat([point_map1,point_map2],dim=2)
conditional_pixel_values_repeat=conditional_pixel_values.repeat(1,14,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=torch.float16)
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)
ref_base_name=os.path.splitext(os.path.basename(ref_image_list[val_id]))[0]
sketch_base_name=os.path.splitext(os.path.basename(each_sample))[0]
supp_dir=os.path.join(args.output_dir,ref_base_name+"_"+sketch_base_name)
os.makedirs(supp_dir, exist_ok=True)
elif args.repeat_sketch:
controlnet_image=[]
for i_2 in range(int(len(control_images)/2)):
sketch=control_images[0]
sketch = np.array(sketch)
sketch=detector(sketch,coarse=False)
sketch=np.repeat(sketch[:, :, np.newaxis], 3, axis=2)
if args.not_quant_sketch:
pass
else:
sketch= (sketch > 200).astype(np.uint8)*255
sketch = Image.fromarray(sketch)
controlnet_image.append(sketch)
for i_3 in range(int(len(control_images)/2)):
sketch=control_images[-1]
sketch = np.array(sketch)
sketch=detector(sketch,coarse=False)
sketch=np.repeat(sketch[:, :, np.newaxis], 3, axis=2)
if args.not_quant_sketch:
pass
else:
sketch= (sketch > 200).astype(np.uint8)*255
sketch = Image.fromarray(sketch)
controlnet_image.append(sketch)
generator = torch.manual_seed(args.seed)
with torch.inference_mode():
video_frames = pipe(
ref_image,
controlnet_condition,
height=args.height,
width=args.width,
num_frames=14,
decode_chunk_size=8,
motion_bucket_id=127,
fps=7,
noise_aug_strength=0.02,
generator=generator,
).frames[0]
out_file = supp_dir+'.mp4'
if args.all_sketch:
export_gif_side_by_side_complete_ablation(ref_image,controlnet_image,video_frames,out_file.replace('.mp4','.gif'),supp_dir,6)
elif args.repeat_sketch:
export_gif_with_ref(control_images[0],video_frames,controlnet_image[-1],controlnet_image[0],out_file.replace('.mp4','.gif'),6)