From 28ccfbfc51068a663a80764e14074df5edf2b5ba Mon Sep 17 00:00:00 2001
From: kongdeqiang <kongdeqiang960204@163.com>
Date: 星期五, 13 三月 2026 17:41:41 +0800
Subject: [PATCH] 提交
---
funasr/auto/auto_model.py | 31 ++++++++++++++++++++++++++-----
1 files changed, 26 insertions(+), 5 deletions(-)
diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index 96c642e..a864dad 100644
--- a/funasr/auto/auto_model.py
+++ b/funasr/auto/auto_model.py
@@ -182,7 +182,10 @@
set_all_random_seed(kwargs.get("seed", 0))
device = kwargs.get("device", "cuda")
- if not torch.cuda.is_available() or kwargs.get("ngpu", 1) == 0:
+ if ((device =="cuda" and not torch.cuda.is_available())
+ or (device == "xpu" and not torch.xpu.is_available())
+ or (device == "mps" and not torch.backends.mps.is_available())
+ or kwargs.get("ngpu", 1) == 0):
device = "cpu"
kwargs["batch_size"] = 1
kwargs["device"] = device
@@ -298,14 +301,27 @@
res = self.model(*args, kwargs)
return res
- def generate(self, input, input_len=None, **cfg):
+ def generate(self, input, input_len=None, progress_callback=None, **cfg):
if self.vad_model is None:
- return self.inference(input, input_len=input_len, **cfg)
+ return self.inference(
+ input, input_len=input_len, progress_callback=progress_callback, **cfg
+ )
else:
- return self.inference_with_vad(input, input_len=input_len, **cfg)
+ return self.inference_with_vad(
+ input, input_len=input_len, progress_callback=progress_callback, **cfg
+ )
- def inference(self, input, input_len=None, model=None, kwargs=None, key=None, **cfg):
+ def inference(
+ self,
+ input,
+ input_len=None,
+ model=None,
+ kwargs=None,
+ key=None,
+ progress_callback=None,
+ **cfg,
+ ):
kwargs = self.kwargs if kwargs is None else kwargs
if "cache" in kwargs:
kwargs.pop("cache")
@@ -362,6 +378,11 @@
if pbar:
pbar.update(end_idx - beg_idx)
pbar.set_description(description)
+ if progress_callback:
+ try:
+ progress_callback(end_idx, num_samples)
+ except Exception as e:
+ logging.error(f"progress_callback error: {e}")
time_speech_total += batch_data_time
time_escape_total += time_escape
--
Gitblit v1.9.1