Update README.md

README.md

@@ -295,7 +295,7 @@ def split_video_into_frames(input_video_path, num_frames, temp_folder='temp_fram
 
     print(f"新的视频已保存到 resampled_video.mp4，包含 {num_frames} 个帧，并保持原始的帧率。")
 
-def generate_video_with_prompt(input_video_path, prompt, model_id, gif_output_path, num_frames=16, keep_imgs=False, temp_folder='temp_frames'):
+def generate_video_with_prompt(input_video_path, prompt, model_id, gif_output_path, seed=0, num_frames=16, keep_imgs=False, temp_folder='temp_frames'):
     """
     生成带有文本提示的视频。
 
@@ -303,6 +303,7 @@ def generate_video_with_prompt(input_video_path, prompt, model_id, gif_output_pa
     :param prompt: 文本提示
     :param model_id: 模型ID
     :param gif_output_path: GIF 输出文件路径
+    :param seed: 随机种子
     :param num_frames: 目标帧数
     :param keep_imgs: 是否保留临时图片
     :param temp_folder: 临时文件夹路径
@@ -319,7 +320,7 @@ def generate_video_with_prompt(input_video_path, prompt, model_id, gif_output_pa
     cn2 = [p2(frame) for frame in conditioning_frames]
 
     negative_prompt = "bad quality, worst quality, jpeg artifacts, ugly"
-    generator = torch.Generator(device="cpu").manual_seed(0)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
 
     pipe = initialize_pipeline(model_id)
 
@@ -352,14 +353,14 @@ if __name__ == "__main__":
     parser.add_argument("prompt", help="文本提示")
     parser.add_argument("model_id", help="模型ID")
     parser.add_argument("gif_output_path", help="GIF 输出文件路径")
+    parser.add_argument("--seed", type=int, default=0, help="随机种子")
     parser.add_argument("--num_frames", type=int, default=16, help="目标帧数")
     parser.add_argument("--keep_imgs", action="store_true", help="是否保留临时图片")
     parser.add_argument("--temp_folder", default='temp_frames', help="临时文件夹路径")
 
     args = parser.parse_args()
 
-    generate_video_with_prompt(args.input_video, args.prompt, args.model_id, args.gif_output_path, args.num_frames, args.keep_imgs, args.temp_folder)
-
+    generate_video_with_prompt(args.input_video, args.prompt, args.model_id, args.gif_output_path, args.seed, args.num_frames, args.keep_imgs, args.temp_folder)
 ```
 
 ```bash