Spaces:
Paused
Paused
File size: 3,562 Bytes
e1ced2b b95bfc8 e1ced2b b95bfc8 e1ced2b b95bfc8 e1ced2b b95bfc8 e1ced2b b95bfc8 e1ced2b b95bfc8 e1ced2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import os
import argparse
from process_pipepline import ProcessPipeline
# simple args container like before
class _Args:
pass
def _parse_args():
args = _Args()
# general paths
args.ckpt_path = "./Wan2.2-Animate-14B/process_checkpoint"
args.video_path = None
args.refer_path = None
args.save_path = None
# processing parameters
args.resolution_area = [1280, 720]
# Keep in sync with wan_engine.FPS. This defaulted to 16 while wan_engine
# defaulted to 8 and NOTHING set the env var, so get_frames built 160 frame
# indices over 80 unique frames — every frame processed twice, then exported
# at 8fps as a 20s half-speed video. wan_engine now force-sets ANIMATE_FPS at
# import so the two can never diverge again; 8 is the matching default.
args.fps = int(os.getenv("ANIMATE_FPS", "8"))
# feature flags
args.replace_flag = True
args.retarget_flag = False
args.use_flux = False
# mask strategy parameters (replacement mode)
args.iterations = 3
args.k = 7
args.w_len = 1
args.h_len = 1
return args
def load_preprocess_models(with_pose=True, pose_device="cuda"):
"""with_pose=False builds a SAM2-only pipeline. Pose now runs on the CPU in
the main process (wan_engine._pose_cpu), so rebuilding the 2.5GB ViTPose ONNX
session inside the paid GPU reservation just to leave it idle is pure waste."""
ckpt_path = "./Wan2.2-Animate-14B/process_checkpoint"
pose2d_checkpoint_path = (
os.path.join(ckpt_path, 'pose2d/vitpose_h_wholebody.onnx') if with_pose else None)
det_checkpoint_path = (
os.path.join(ckpt_path, 'det/yolov10m.onnx') if with_pose else None)
sam2_checkpoint_path = [os.path.join(ckpt_path, 'sam2/sam2_hiera_large.pt'),"sam2_hiera_l.yaml"]
flux_kontext_path = 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, pose_device=pose_device)
return process_pipeline
def run(process_pipeline, input_video, edited_frame, preprocess_dir, w, h, tag_string,
pts_by_frame: dict, lbs_by_frame: dict,
frames=None, precomputed_metas=None):
args = _parse_args()
if tag_string == "retarget_flag":
retarget_flag = True
replace_flag = False
else:
retarget_flag = False
replace_flag = True
src_pose_path, src_face_path, src_bg_path, src_mask_path, src_ref_image = process_pipeline(video_path=input_video,
refer_image_path=edited_frame,
output_path=preprocess_dir,
resolution_area=[w, h],
fps=args.fps,
iterations=args.iterations,
k=args.k,
w_len=args.w_len,
h_len=args.h_len,
retarget_flag=retarget_flag,
use_flux=args.use_flux,
replace_flag=replace_flag,
pts_by_frame=pts_by_frame,
lbs_by_frame=lbs_by_frame,
frames=frames,
precomputed_metas=precomputed_metas)
return src_pose_path, src_face_path, src_bg_path, src_mask_path, src_ref_image
|