Update inference_from_video.py

inference_from_video.py

@@ -66,15 +66,15 @@ def parse_args():
     )
     parser.add_argument(
         "--sample_rate", type=int, default=-1,
-        help="How many test instances to evaluate.",
+        help="Sample rate for audio output."
     )
     parser.add_argument(
         "--save_dir", type=str, default="./outputs/tmp",
-        help="output save dir"
+        help="Output save directory"
     )
     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 for item in wave]
+            all_outputs += [item.cpu().numpy() for item in wave]  # Ensure wave is on CPU and in numpy format
             
     # Save #
     exp_id = str(int(time.time()))
@@ -194,7 +194,12 @@ 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):
-            sf.write("{}/{}".format(output_dir, wavname[j]), wav, samplerate=sample_rate)
+            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}")
             
     else:
         for i in range(num_samples):
@@ -208,8 +213,12 @@ def main():
             ranked_wavs_for_text = [wavs_for_text[r] for r in rank]
             
             for i, wav in enumerate(ranked_wavs_for_text):
-                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)
+                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}")
             
 if __name__ == "__main__":
     main()