scripts_infer/anidoc_inference.py

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)