update

app.py

@@ -132,7 +132,8 @@ laionclap_model = load_laionclap()
 model = prepare_model(
     model_config=model_config, 
     clap_config=clap_config, 
-    checkpoint_path='chat.pt'
+    checkpoint_path='chat.pt',
+    device=device
 )
 
 
@@ -147,6 +148,7 @@ def inference_item(name, prompt):
     outputs = inference(
         model, text_tokenizer, item, processed_item,
         inference_kwargs,
+        device=device
     )
 
     laionclap_scores = compute_laionclap_text_audio_sim(

inference_utils.py

@@ -33,7 +33,7 @@ def prepare_tokenizer(model_config):
     return text_tokenizer
 
 
-def prepare_model(model_config, clap_config, checkpoint_path, device_id=0):
+def prepare_model(model_config, clap_config, checkpoint_path, device=0):
     os.environ["TOKENIZERS_PARALLELISM"] = "false"  # disable the tokenizer parallelism warning
     model, tokenizer = create_model_and_transforms(
         **model_config,
@@ -43,7 +43,7 @@ def prepare_model(model_config, clap_config, checkpoint_path, device_id=0):
         freeze_lm_embeddings=False,
     )
     model.eval()
-    model = model.to(device_id)
+    model = model.to(device)
 
     checkpoint = torch.load(checkpoint_path, map_location="cpu")
     model_state_dict = checkpoint["model_state_dict"]
@@ -53,11 +53,11 @@ def prepare_model(model_config, clap_config, checkpoint_path, device_id=0):
     return model
 
 
-def inference(model, tokenizer, item, processed_item, inference_kwargs, device_id=0):
+def inference(model, tokenizer, item, processed_item, inference_kwargs, device=0):
     filename, audio_clips, audio_embed_mask, input_ids, attention_mask = processed_item
-    audio_clips = audio_clips.to(device_id, dtype=None, non_blocking=True)
-    audio_embed_mask = audio_embed_mask.to(device_id, dtype=None, non_blocking=True)
-    input_ids = input_ids.to(device_id, dtype=None, non_blocking=True).squeeze()
+    audio_clips = audio_clips.to(device, dtype=None, non_blocking=True)
+    audio_embed_mask = audio_embed_mask.to(device, dtype=None, non_blocking=True)
+    input_ids = input_ids.to(device, dtype=None, non_blocking=True).squeeze()
 
     media_token_id = tokenizer.encode("<audio>")[-1]
     eoc_token_id = tokenizer.encode("<|endofchunk|>")[-1]