update space

app.py

@@ -73,10 +73,11 @@ beats_path = hf_hub_download(
 
 model_path = "THUdyh/Ola-7b"
 tokenizer, model, image_processor, _ = load_pretrained_model(model_path, None)
-model = model.to('cuda').eval()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device).eval()
 model = model.bfloat16()
 
-tts_model = CosyVoice('iic/CosyVoice-300M-SFT', load_jit=True, fp16=True)
+# tts_model = CosyVoice('iic/CosyVoice-300M-SFT', load_jit=True, fp16=True)
 # tts_model = CosyVoice('FunAudioLLM/CosyVoice-300M-SFT', load_jit=True, fp16=True)
 OUTPUT_SPEECH = False
 
@@ -186,10 +187,10 @@ def ola_inference(multimodal, audio_path):
     if USE_SPEECH and audio_path:
         audio_path = audio_path
         speech, speech_length, speech_chunk, speech_wav = load_audio(audio_path)
-        speechs.append(speech.bfloat16().to('cuda'))
-        speech_lengths.append(speech_length.to('cuda'))
-        speech_chunks.append(speech_chunk.to('cuda'))
-        speech_wavs.append(speech_wav.to('cuda'))
+        speechs.append(speech.bfloat16().to(device))
+        speech_lengths.append(speech_length.to(device))
+        speech_chunks.append(speech_chunk.to(device))
+        speech_wavs.append(speech_wav.to(device))
         print('load audio')
     elif USE_SPEECH and not audio_path:
         # parse audio in the video
@@ -197,15 +198,15 @@ def ola_inference(multimodal, audio_path):
         audio.write_audiofile("./video_audio.wav")
         video_audio_path = './video_audio.wav'
         speech, speech_length, speech_chunk, speech_wav = load_audio(video_audio_path)
-        speechs.append(speech.bfloat16().to('cuda'))
-        speech_lengths.append(speech_length.to('cuda'))
-        speech_chunks.append(speech_chunk.to('cuda'))
-        speech_wavs.append(speech_wav.to('cuda'))
+        speechs.append(speech.bfloat16().to(device))
+        speech_lengths.append(speech_length.to(device))
+        speech_chunks.append(speech_chunk.to(device))
+        speech_wavs.append(speech_wav.to(device))
     else:
-        speechs = [torch.zeros(1, 3000, 128).bfloat16().to('cuda')]
-        speech_lengths = [torch.LongTensor([3000]).to('cuda')]
-        speech_wavs = [torch.zeros([1, 480000]).to('cuda')]
-        speech_chunks = [torch.LongTensor([1]).to('cuda')]
+        speechs = [torch.zeros(1, 3000, 128).bfloat16().to(device)]
+        speech_lengths = [torch.LongTensor([3000]).to(device)]
+        speech_wavs = [torch.zeros([1, 480000]).to(device)]
+        speech_chunks = [torch.LongTensor([1]).to(device)]
     
     conv_mode = "qwen_1_5"
     if text:
@@ -224,11 +225,11 @@ def ola_inference(multimodal, audio_path):
     conv.append_message(conv.roles[1], None)
     prompt = conv.get_prompt()
     if USE_SPEECH and audio_path:
-        input_ids = tokenizer_speech_question_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to('cuda')
+        input_ids = tokenizer_speech_question_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(device)
     elif USE_SPEECH:
-        input_ids = tokenizer_speech_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to('cuda')
+        input_ids = tokenizer_speech_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(device)
     else:
-        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to('cuda')
+        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(device)
 
     if modality == "video":
         video_processed = []
@@ -272,7 +273,7 @@ def ola_inference(multimodal, audio_path):
 
     pad_token_ids = 151643
 
-    attention_masks = input_ids.ne(pad_token_ids).long().to('cuda')
+    attention_masks = input_ids.ne(pad_token_ids).long().to(device)
     stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
     keywords = [stop_str]
     stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)