Update inference_from_video.py

inference_from_video.py

@@ -3,16 +3,10 @@ import copy
 import json
 import time
 import torch
-
-# Check if CUDA is available and set the device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print("Using device:", device)
-
 import argparse
 from PIL import Image
 import numpy as np
 import soundfile as sf
-#import wandb
 from tqdm import tqdm
 from diffusers import DDPMScheduler
 from models import build_pretrained_models, AudioDiffusion
@@ -21,6 +15,10 @@ import torchaudio
 import tools.torch_tools as torch_tools
 from datasets import load_dataset
 
+# Check if CUDA is available and set the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+
 class dotdict(dict):
     """dot.notation access to dictionary attributes"""
     __getattr__ = dict.get
@@ -80,7 +78,6 @@ def parse_args():
     )
     
     args = parser.parse_args()
-
     return args
 
 def main():
@@ -93,12 +90,13 @@ def main():
     # Load Models #
     name = train_args.vae_model
     vae, stft = build_pretrained_models(name)
-    vae, stft = vae.cuda(), stft.cuda()
+    vae, stft = vae.to(device), stft.to(device)  # Ensure models are on the correct device
+    
     model_class = AudioDiffusion
     if train_args.ib:
         print("*****USING MODEL IMAGEBIND*****")
         from models_imagebind import AudioDiffusion_IB
-        model_class = AudioDiffusion if not train_args.ib else AudioDiffusion_IB
+        model_class = AudioDiffusion_IB
     elif train_args.lb:
         print("*****USING MODEL LANGUAGEBIND*****")
         from models_languagebind import AudioDiffusion_LB
@@ -125,9 +123,8 @@ def main():
     model.eval()
 
     # Load Trained Weight #
-    
     if args.model.endswith(".pt") or args.model.endswith(".bin"):
-        model.load_state_dict(torch.load(args.model), strict=False)
+        model.load_state_dict(torch.load(args.model, map_location=device), strict=False)
     else:
         from safetensors.torch import load_model
         load_model(model, args.model, strict=False)
@@ -136,8 +133,6 @@ def main():
     
     scheduler = DDPMScheduler.from_pretrained(train_args.scheduler_name, subfolder="scheduler")
     sample_rate = args.sample_rate
-    #evaluator = EvaluationHelper(16000, "cuda:0")
-    
 
     def audio_text_matching(waveforms, text, sample_freq=24000, max_len_in_seconds=10):
         new_freq = 48000
@@ -163,7 +158,6 @@ def main():
     else:
         prefix = ""
 
-    # data_path = "data/video_test/"
     data_path = args.data_path
     wavname = [f"{name.split('.')[0]}.wav" for name in os.listdir(data_path)]
     video_features = []
@@ -171,25 +165,15 @@ def main():
         video_path = os.path.join(data_path, video_file)
         video_feature = torch_tools.load_video(video_path, frame_rate=2, size=224)
         print(video_feature.shape)
-        video_features.append(video_feature)
+        video_features.append(video_feature.to(device))  # Move to device
     
     # Generate #
     num_steps, guidance, batch_size, num_samples = args.num_steps, args.guidance, args.batch_size, args.num_samples
     all_outputs = []
         
     for k in tqdm(range(0, len(wavname), batch_size)):
-        
         with torch.no_grad():
-            # if train_args.task == 'image2audio':
-            #     prompt = text_prompts[k: k+batch_size]
-            #     imgs = []
-            #     for img_path in prompt:
-            #         img = Image.open(img_path)
-            #         imgs.append(np.array(img))
-            #     prompt = imgs
-            # elif train_args.task == 'video2audio':
             prompt = video_features[k: k+batch_size]
-
             latents = model.inference(scheduler, None, prompt, None, num_steps, guidance, num_samples, disable_progress=True, device=device)
             mel = vae.decode_first_stage(latents)
             wave = vae.decode_to_waveform(mel)