Update inference_from_video.py

inference_from_video.py

@@ -65,16 +65,16 @@ def parse_args():
         help="How many test instances to evaluate.",
     )
     parser.add_argument(
-        "--sample_rate", type=int, default=-1,
-        help="Sample rate for audio output."
+        "--sample_rate", type=int, default=48000,
+        help="Sample rate for audio output.",
     )
     parser.add_argument(
         "--save_dir", type=str, default="./outputs/tmp",
-        help="Output save directory"
+        help="output save dir"
     )
     parser.add_argument(
         "--data_path", type=str, default="data/video_processed/video_gt_augment",
-        help="Inference data path"
+        help="inference data path"
     )
     
     args = parser.parse_args()
@@ -183,7 +183,7 @@ def main():
             mel = vae.decode_first_stage(latents)
             wave = vae.decode_to_waveform(mel)
             
-            all_outputs += [item.cpu().numpy() for item in wave]  # Ensure wave is on CPU and in numpy format
+            all_outputs += [item for item in wave]
             
     # Save #
     exp_id = str(int(time.time()))
@@ -194,12 +194,7 @@ def main():
         output_dir = "{}/{}_{}_steps_{}_guidance_{}_sampleRate_{}_augment".format(args.save_dir, exp_id, "_".join(args.model.split("/")[1:-1]), num_steps, guidance, sample_rate)
         os.makedirs(output_dir, exist_ok=True)
         for j, wav in enumerate(all_outputs):
-            file_path = "{}/{}".format(output_dir, wavname[j])
-            try:
-                sf.write(file_path, wav, samplerate=sample_rate)
-                print(f"Saved {file_path}")
-            except Exception as e:
-                print(f"Error saving {file_path}: {e}")
+            sf.write("{}/{}".format(output_dir, wavname[j]), wav, samplerate=sample_rate)
             
     else:
         for i in range(num_samples):
@@ -213,12 +208,8 @@ def main():
             ranked_wavs_for_text = [wavs_for_text[r] for r in rank]
             
             for i, wav in enumerate(ranked_wavs_for_text):
-                file_path = "{}/{}".format(output_dir, wavname[k])
-                try:
-                    sf.write(file_path, wav, samplerate=sample_rate)
-                    print(f"Saved {file_path}")
-                except Exception as e:
-                    print(f"Error saving {file_path}: {e}")
+                output_dir = "{}/{}_{}_steps_{}_guidance_{}_sampleRate_{}/rank_{}".format(args.save_dir, exp_id, "_".join(args.model.split("/")[1:-1]), num_steps, guidance, sample_rate, i+1)
+                sf.write("{}/{}".format(output_dir, wavname[k]), wav, samplerate=sample_rate)
             
 if __name__ == "__main__":
-    main()
+    main()