From 713d69f3902455fcc182672cfd58e588c7022944 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 20 Jan 2025 02:26:37 -0500 Subject: [PATCH] accept pose videos --- .../StableAnimator/inference_advanced.py | 154 +++++++++++++++--- 1 file changed, 131 insertions(+), 23 deletions(-) diff --git a/animation/StableAnimator/inference_advanced.py b/animation/StableAnimator/inference_advanced.py index fe3317c..a0ff0fa 100644 --- a/animation/StableAnimator/inference_advanced.py +++ b/animation/StableAnimator/inference_advanced.py @@ -16,6 +16,8 @@ from animation.modules.pose_net import PoseNet from animation.modules.unet import UNetSpatioTemporalConditionModel from animation.pipelines.inference_pipeline_animation import InferenceAnimationPipeline import random +import subprocess +from pathlib import Path def seed_everything(seed): torch.manual_seed(seed) @@ -78,9 +80,10 @@ def export_to_gif(frames, output_gif_path, fps): def parse_args(): parser = argparse.ArgumentParser( - description="Script to train Stable Diffusion XL for InstructPix2Pix." + description="Inference script" ) + # KEEP parser.add_argument( "--pretrained_model_name_or_path", type=str, @@ -88,8 +91,20 @@ def parse_args(): required=True ) + ## KEEP + #parser.add_argument( + # "--validation_image", + # type=str, + # default=None, + # help=( + # "A set of paths to the controlnext conditioning image be evaluated every `--validation_steps`" + # " and logged to `--report_to`. Provide either a matching number of `--validation_prompt`s, a" + # " a single `--validation_prompt` to be used with all `--validation_image`s, or a single" + # " `--validation_image` that will be used with all `--validation_prompt`s." + # ), + #) parser.add_argument( - "--validation_image", + "--start_image_path", type=str, default=None, help=( @@ -99,15 +114,41 @@ def parse_args(): " `--validation_image` that will be used with all `--validation_prompt`s." ), ) + + # =================== + # Pose Input + # Supply either `pose_images_folder`, `pose_video_path`, or `pre_pose_video_path`. + # - pose_images_folder has inputs already fully processed + # - pose_video_path just needs to be converted to frames + # - pre_pose_video_path has not been converted to pose data frames + # =================== parser.add_argument( - "--validation_control_folder", + "--pose_images_folder", type=str, default=None, help=( - "the validation control image" + "the validation control images" + ), + ) + parser.add_argument( + "--pose_video_path", + type=str, + default=None, + help=( + "the validation control video (optional)" + ), + ) + parser.add_argument( + "--pre_pose_video_path", + type=str, + default=None, + help=( + "video to convert to pose" ), ) + + # TODO FIX parser.add_argument( "--output_dir", type=str, @@ -115,6 +156,7 @@ def parse_args(): required=True ) + # TODO FIX parser.add_argument( "--height", type=int, @@ -122,6 +164,7 @@ def parse_args(): required=False ) + # TODO FIX parser.add_argument( "--width", type=int, @@ -129,6 +172,7 @@ def parse_args(): required=False ) + # KEEP parser.add_argument( "--guidance_scale", type=float, @@ -136,6 +180,7 @@ def parse_args(): required=False ) + # KEEP parser.add_argument( "--num_inference_steps", type=int, @@ -143,18 +188,23 @@ def parse_args(): required=False ) + # KEEP parser.add_argument( "--posenet_model_name_or_path", type=str, default=None, help="Path to pretrained posenet model", ) + + # KEEP parser.add_argument( "--face_encoder_model_name_or_path", type=str, default=None, help="Path to pretrained face encoder model", ) + + # KEEP parser.add_argument( "--unet_model_name_or_path", type=str, @@ -162,6 +212,7 @@ def parse_args(): help="Path to pretrained unet model", ) + # KEEP parser.add_argument( "--tile_size", type=int, @@ -169,6 +220,7 @@ def parse_args(): required=False ) + # KEEP parser.add_argument( "--overlap", type=int, @@ -176,23 +228,30 @@ def parse_args(): required=False ) + # KEEP parser.add_argument( "--noise_aug_strength", type=float, default=0.0, # or set to 0.02 required=False ) + + # KEEP parser.add_argument( "--frames_overlap", type=int, default=4, required=False ) + + # KEEP parser.add_argument( "--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", ) + + # KEEP parser.add_argument( "--revision", type=str, @@ -200,6 +259,8 @@ def parse_args(): required=False, help="Revision of pretrained model identifier from huggingface.co/models.", ) + + # KEEP parser.add_argument( "--decode_chunk_size", type=int, @@ -210,6 +271,52 @@ def parse_args(): args = parser.parse_args() return args +def split_video_to_frames(video_path, output_path): + Path(output_path).mkdir(parents=True, exist_ok=True) + filename_format = f"{output_path}/frame_%d.png" + command = [ + 'ffmpeg', + '-i', video_path, + '-q:v', '1', + '-start_number', '0', + filename_format, + ] + print(f"Command: {command}", flush=True) + result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True) + #video_info = json.loads(result.stdout) + +def prepare_pose_frames(args): + pose_images_dir = args.pose_images_folder + + if not pose_images_dir: + raise Exception('pose_images_folder must be set, even when using other arguments') + + Path(pose_images_dir).mkdir(parents=True, exist_ok=True) + + if args.pre_pose_video_path: + print("Preparing pose frames from pre-pose video.") + pre_pose_frame_dir = Path(f"{pose_images_dir}/frames") + pre_pose_frame_dir.mkdir(parents=True, exist_ok=True) + split_video_to_frames(args.pre_pose_video_path, pre_pose_frame_dir) + reference_image_path = pre_pose_frame_dir / "frame_0.png" + command = [ + "python", "DWPose/skeleton_extraction.py", + "--target_image_folder_path", pre_pose_frame_dir, + "--ref_image_path", reference_image_path, + "--poses_folder_path", pose_images_dir, + ] + print(f"Command: {command}", flush=True) + result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True) + + elif args.pose_video_path: + print("Preparing pose frames from pose video.") + split_video_to_frames(args.pose_video_path, pose_images_dir) + else: + print("Pose frames assumed to already exist.") + + # TODO: autodetect resolution + pose_images = load_images_from_folder(pose_images_dir, width=args.width, height=args.height) + return pose_images if __name__ == "__main__": args = parse_args() @@ -333,39 +440,40 @@ if __name__ == "__main__": os.makedirs(args.output_dir, exist_ok=True) - validation_image_path = args.validation_image - validation_image = Image.open(args.validation_image).convert('RGB') - validation_control_images = load_images_from_folder(args.validation_control_folder, width=args.width, height=args.height) + pose_images = prepare_pose_frames(args) + num_frames = len(pose_images) + + start_image_path = args.start_image_path + start_image = Image.open(args.start_image_path).convert('RGB') - num_frames = len(validation_control_images) face_model.face_helper.clean_all() - validation_face = cv2.imread(validation_image_path) - validation_image_bgr = cv2.cvtColor(validation_face, cv2.COLOR_RGB2BGR) - validation_image_face_info = face_model.app.get(validation_image_bgr) - if len(validation_image_face_info) > 0: - validation_image_face_info = sorted(validation_image_face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1] - validation_image_id_ante_embedding = validation_image_face_info['embedding'] + start_image_face = cv2.imread(start_image_path) + start_image_bgr = cv2.cvtColor(start_image_face, cv2.COLOR_RGB2BGR) + start_image_face_info = face_model.app.get(start_image_bgr) + if len(start_image_face_info) > 0: + start_image_face_info = sorted(start_image_face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1] + start_image_id_ante_embedding = start_image_face_info['embedding'] else: - validation_image_id_ante_embedding = None + start_image_id_ante_embedding = None - if validation_image_id_ante_embedding is None: - face_model.face_helper.read_image(validation_image_bgr) + if start_image_id_ante_embedding is None: + face_model.face_helper.read_image(start_image_bgr) face_model.face_helper.get_face_landmarks_5(only_center_face=True) face_model.face_helper.align_warp_face() if len(face_model.face_helper.cropped_faces) == 0: - validation_image_id_ante_embedding = np.zeros((512,)) + start_image_id_ante_embedding = np.zeros((512,)) else: - validation_image_align_face = face_model.face_helper.cropped_faces[0] + start_image_align_face = face_model.face_helper.cropped_faces[0] print('fail to detect face using insightface, extract embedding on align face') - validation_image_id_ante_embedding = face_model.handler_ante.get_feat(validation_image_align_face) + start_image_id_ante_embedding = face_model.handler_ante.get_feat(start_image_align_face) # generator = torch.Generator(device=accelerator.device).manual_seed(23123134) decode_chunk_size = args.decode_chunk_size video_frames = pipeline( - image=validation_image, - image_pose=validation_control_images, + image=start_image, + image_pose=pose_images, height=args.height, width=args.width, num_frames=num_frames, @@ -380,7 +488,7 @@ if __name__ == "__main__": num_inference_steps=args.num_inference_steps, generator=generator, output_type="pil", - validation_image_id_ante_embedding=validation_image_id_ante_embedding, + validation_image_id_ante_embedding=start_image_id_ante_embedding, ).frames[0] out_file = os.path.join(