feat:增加device 参数

fix: set_kv_cache使用默认device问题
This commit is contained in:
wuhongsheng 2024-09-06 10:56:48 +08:00
parent ae04f97733
commit a97e7cd41a
2 changed files with 3 additions and 3 deletions

View File

@ -399,7 +399,7 @@ class OmniInference:
model = self.model model = self.model
with self.fabric.init_tensor(): with self.fabric.init_tensor():
model.set_kv_cache(batch_size=2) model.set_kv_cache(batch_size=2,device=self.device)
mel, leng = load_audio(audio_path) mel, leng = load_audio(audio_path)
audio_feature, input_ids = get_input_ids_whisper_ATBatch(mel, leng, self.whispermodel, self.device) audio_feature, input_ids = get_input_ids_whisper_ATBatch(mel, leng, self.whispermodel, self.device)

View File

@ -46,9 +46,9 @@ def create_app():
return server.server return server.server
def serve(ip='0.0.0.0', port=60808): def serve(ip='0.0.0.0', port=60808, device='cuda:0'):
OmniChatServer(ip, port=port, run_app=True) OmniChatServer(ip, port=port,run_app=True, device=device)
if __name__ == "__main__": if __name__ == "__main__":