|
|
|
|
|
import os |
|
|
import argparse |
|
|
from process_pipepline import ProcessPipeline |
|
|
|
|
|
|
|
|
def _parse_args(): |
|
|
parser = argparse.ArgumentParser( |
|
|
description="The preprocessing pipeline for Wan-animate." |
|
|
) |
|
|
|
|
|
parser.add_argument( |
|
|
"--ckpt_path", |
|
|
type=str, |
|
|
default=None, |
|
|
help="The path to the preprocessing model's checkpoint directory. ") |
|
|
|
|
|
parser.add_argument( |
|
|
"--video_path", |
|
|
type=str, |
|
|
default=None, |
|
|
help="The path to the driving video.") |
|
|
parser.add_argument( |
|
|
"--refer_path", |
|
|
type=str, |
|
|
default=None, |
|
|
help="The path to the refererence image.") |
|
|
parser.add_argument( |
|
|
"--save_path", |
|
|
type=str, |
|
|
default=None, |
|
|
help="The path to save the processed results.") |
|
|
|
|
|
parser.add_argument( |
|
|
"--resolution_area", |
|
|
type=int, |
|
|
nargs=2, |
|
|
default=[1280, 720], |
|
|
help="The target resolution for processing, specified as [width, height]. To handle different aspect ratios, the video is resized to have a total area equivalent to width * height, while preserving the original aspect ratio." |
|
|
) |
|
|
parser.add_argument( |
|
|
"--fps", |
|
|
type=int, |
|
|
default=30, |
|
|
help="The target FPS for processing the driving video. Set to -1 to use the video's original FPS." |
|
|
) |
|
|
|
|
|
parser.add_argument( |
|
|
"--replace_flag", |
|
|
action="store_true", |
|
|
default=False, |
|
|
help="Whether to use replacement mode.") |
|
|
parser.add_argument( |
|
|
"--retarget_flag", |
|
|
action="store_true", |
|
|
default=False, |
|
|
help="Whether to use pose retargeting. Currently only supported in animation mode") |
|
|
parser.add_argument( |
|
|
"--use_flux", |
|
|
action="store_true", |
|
|
default=False, |
|
|
help="Whether to use image editing in pose retargeting. Recommended if the character in the reference image or the first frame of the driving video is not in a standard, front-facing pose") |
|
|
|
|
|
|
|
|
parser.add_argument( |
|
|
"--iterations", |
|
|
type=int, |
|
|
default=3, |
|
|
help="Number of iterations for mask dilation." |
|
|
) |
|
|
parser.add_argument( |
|
|
"--k", |
|
|
type=int, |
|
|
default=7, |
|
|
help="Number of kernel size for mask dilation." |
|
|
) |
|
|
parser.add_argument( |
|
|
"--w_len", |
|
|
type=int, |
|
|
default=1, |
|
|
help="The number of subdivisions for the grid along the 'w' dimension. A higher value results in a more detailed contour. A value of 1 means no subdivision is performed." |
|
|
) |
|
|
parser.add_argument( |
|
|
"--h_len", |
|
|
type=int, |
|
|
default=1, |
|
|
help="The number of subdivisions for the grid along the 'h' dimension. A higher value results in a more detailed contour. A value of 1 means no subdivision is performed." |
|
|
) |
|
|
args = parser.parse_args() |
|
|
|
|
|
return args |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
args = _parse_args() |
|
|
args_dict = vars(args) |
|
|
print(args_dict) |
|
|
|
|
|
assert len(args.resolution_area) == 2, "resolution_area should be a list of two integers [width, height]" |
|
|
assert not args.use_flux or args.retarget_flag, "Image editing with FLUX can only be used when pose retargeting is enabled." |
|
|
|
|
|
pose2d_checkpoint_path = os.path.join(args.ckpt_path, 'pose2d/vitpose_h_wholebody.onnx') |
|
|
det_checkpoint_path = os.path.join(args.ckpt_path, 'det/yolov10m.onnx') |
|
|
|
|
|
sam2_checkpoint_path = os.path.join(args.ckpt_path, 'sam2/sam2_hiera_large.pt') if args.replace_flag else None |
|
|
flux_kontext_path = os.path.join(args.ckpt_path, 'FLUX.1-Kontext-dev') if args.use_flux else None |
|
|
process_pipeline = ProcessPipeline(det_checkpoint_path=det_checkpoint_path, pose2d_checkpoint_path=pose2d_checkpoint_path, sam_checkpoint_path=sam2_checkpoint_path, flux_kontext_path=flux_kontext_path) |
|
|
os.makedirs(args.save_path, exist_ok=True) |
|
|
process_pipeline(video_path=args.video_path, |
|
|
refer_image_path=args.refer_path, |
|
|
output_path=args.save_path, |
|
|
resolution_area=args.resolution_area, |
|
|
fps=args.fps, |
|
|
iterations=args.iterations, |
|
|
k=args.k, |
|
|
w_len=args.w_len, |
|
|
h_len=args.h_len, |
|
|
retarget_flag=args.retarget_flag, |
|
|
use_flux=args.use_flux, |
|
|
replace_flag=args.replace_flag) |
|
|
|
|
|
|
