Update FantasyTalking/infer.py

FantasyTalking/infer.py

@@ -18,25 +18,112 @@ from FantasyTalking.utils import get_audio_features, resize_image_by_longest_edg
 
 
 def parse_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()
+    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
 
 
 def load_models(args):
@@ -61,7 +148,9 @@ 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")
@@ -86,7 +175,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}")
@@ -128,7 +217,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=latents_num_frames,
+        latents_num_frames=(num_frames - 1) // 4 + 1,
     )
     print("✅ Video frames generated.")
 
@@ -158,4 +247,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}")