Update FantasyTalking/infer.py

FantasyTalking/infer.py

@@ -18,112 +18,25 @@ from FantasyTalking.utils import get_audio_features, resize_image_by_longest_edg
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description="Simple example of a training script.")
-
-    parser.add_argument(
-        "--wan_model_dir",
-        type=str,
-        default="./models/Wan2.1-I2V-14B-720P",
-        required=False,
-        help="The dir of the Wan I2V 14B model.",
-    )
-    parser.add_argument(
-        "--fantasytalking_model_path",
-        type=str,
-        default="./models/fantasytalking_model.ckpt",
-        required=False,
-        help="The .ckpt path of fantasytalking model.",
-    )
-    parser.add_argument(
-        "--wav2vec_model_dir",
-        type=str,
-        default="./models/wav2vec2-base-960h",
-        required=False,
-        help="The dir of wav2vec model.",
-    )
-
-    parser.add_argument(
-        "--image_path",
-        type=str,
-        default="./assets/images/woman.png",
-        required=False,
-        help="The path of the image.",
-    )
-
-    parser.add_argument(
-        "--audio_path",
-        type=str,
-        default="./assets/audios/woman.wav",
-        required=False,
-        help="The path of the audio.",
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        default="A woman is talking.",
-        required=False,
-        help="prompt.",
-    )
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default="./output",
-        help="Dir to save the model.",
-    )
-    parser.add_argument(
-        "--image_size",
-        type=int,
-        default=512,
-        help="The image will be resized proportionally to this size.",
-    )
-    parser.add_argument(
-        "--audio_scale",
-        type=float,
-        default=1.0,
-        help="Audio condition injection weight",
-    )
-    parser.add_argument(
-        "--prompt_cfg_scale",
-        type=float,
-        default=5.0,
-        required=False,
-        help="Prompt cfg scale",
-    )
-    parser.add_argument(
-        "--audio_cfg_scale",
-        type=float,
-        default=5.0,
-        required=False,
-        help="Audio cfg scale",
-    )
-    parser.add_argument(
-        "--max_num_frames",
-        type=int,
-        default=81,
-        required=False,
-        help="The maximum frames for generating videos, the audio part exceeding max_num_frames/fps will be truncated.",
-    )
-    parser.add_argument(
-        "--fps",
-        type=int,
-        default=23,
-        required=False,
-    )
-    parser.add_argument(
-        "--num_persistent_param_in_dit",
-        type=int,
-        default=None,
-        required=False,
-        help="Maximum parameter quantity retained in video memory, small number to reduce VRAM required",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=1111,
-        required=False,
-    )
-    args = parser.parse_args()
-    return args
+    parser = argparse.ArgumentParser(description="FantasyTalking Video Generator")
+
+    parser.add_argument("--wan_model_dir", type=str, default="./models/Wan2.1-I2V-14B-720P")
+    parser.add_argument("--fantasytalking_model_path", type=str, default="./models/fantasytalking_model.ckpt")
+    parser.add_argument("--wav2vec_model_dir", type=str, default="./models/wav2vec2-base-960h")
+    parser.add_argument("--image_path", type=str, default="./assets/images/woman.png")
+    parser.add_argument("--audio_path", type=str, default="./assets/audios/woman.wav")
+    parser.add_argument("--prompt", type=str, default="A woman is talking.")
+    parser.add_argument("--output_dir", type=str, default="./output")
+    parser.add_argument("--image_size", type=int, default=512)
+    parser.add_argument("--audio_scale", type=float, default=1.0)
+    parser.add_argument("--prompt_cfg_scale", type=float, default=5.0)
+    parser.add_argument("--audio_cfg_scale", type=float, default=5.0)
+    parser.add_argument("--max_num_frames", type=int, default=81)
+    parser.add_argument("--fps", type=int, default=23)
+    parser.add_argument("--num_persistent_param_in_dit", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=1111)
+
+    return parser.parse_args()
 
 
 def load_models(args):
@@ -148,9 +61,7 @@ def load_models(args):
     )
     print("✅ Wan I2V models loaded.")
 
-    pipe = WanVideoPipeline.from_model_manager(
-        model_manager, torch_dtype=torch.bfloat16, device="cuda"
-    )
+    pipe = WanVideoPipeline.from_model_manager(model_manager, torch_dtype=torch.bfloat16, device="cuda")
 
     print("🔄 Loading FantasyTalking model...")
     fantasytalking = FantasyTalkingAudioConditionModel(pipe.dit, 768, 2048).to("cuda")
@@ -175,7 +86,7 @@ def main(args, pipe, fantasytalking, wav2vec_processor, wav2vec):
     print(f"🔊 Getting duration of audio: {args.audio_path}")
     duration = librosa.get_duration(filename=args.audio_path)
     print(f"🎞️ Duration: {duration:.2f}s")
-   
+
     latents_num_frames = min(int(duration * args.fps / 4), args.max_num_frames // 4)
     num_frames = (latents_num_frames - 1) * 4
     print(f"📽️ Calculated number of frames: {num_frames}")
@@ -217,7 +128,7 @@ def main(args, pipe, fantasytalking, wav2vec_processor, wav2vec):
         audio_cfg_scale=args.audio_cfg_scale,
         audio_proj=audio_proj_split,
         audio_context_lens=audio_context_lens,
-        latents_num_frames=(num_frames - 1) // 4 + 1,
+        latents_num_frames=latents_num_frames,
     )
     print("✅ Video frames generated.")
 
@@ -247,4 +158,4 @@ if __name__ == "__main__":
     args = parse_args()
     pipe, fantasytalking, wav2vec_processor, wav2vec = load_models(args)
     video_path = main(args, pipe, fantasytalking, wav2vec_processor, wav2vec)
-    print(f"🎉 Done! Final video path: {video_path}")
+    print(f"🎉 Done! Final video path: {video_path}")