Support loading hf models

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>
This commit is contained in:
makaveli10
2024-11-26 11:14:18 +05:30
parent 446fc6e835
commit 2eff360b9e
2 changed files with 41 additions and 30 deletions
+1 -1
View File
@@ -87,7 +87,7 @@ client = TranscriptionClient(
9090,
lang="en",
translate=False,
model="small",
model="small", # also support hf_model => `Systran/faster-whisper-small`
use_vad=False,
save_output_recording=True, # Only used for microphone input, False by Default
output_recording_filename="./output_recording.wav", # Only used for microphone input
+40 -29
View File
@@ -181,22 +181,26 @@ class TranscriptionServer:
}))
self.backend = BackendType.FASTER_WHISPER
if self.backend.is_faster_whisper():
if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
logging.info(f"Using custom model {faster_whisper_custom_model_path}")
options["model"] = faster_whisper_custom_model_path
client = ServeClientFasterWhisper(
websocket,
language=options["language"],
task=options["task"],
client_uid=options["uid"],
model=options["model"],
initial_prompt=options.get("initial_prompt"),
vad_parameters=options.get("vad_parameters"),
use_vad=self.use_vad,
single_model=self.single_model,
)
logging.info("Running faster_whisper backend.")
try:
if self.backend.is_faster_whisper():
if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
logging.info(f"Using custom model {faster_whisper_custom_model_path}")
options["model"] = faster_whisper_custom_model_path
client = ServeClientFasterWhisper(
websocket,
language=options["language"],
task=options["task"],
client_uid=options["uid"],
model=options["model"],
initial_prompt=options.get("initial_prompt"),
vad_parameters=options.get("vad_parameters"),
use_vad=self.use_vad,
single_model=self.single_model,
)
logging.info("Running faster_whisper backend.")
except Exception as e:
return
if client is None:
raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
@@ -224,7 +228,7 @@ class TranscriptionServer:
logging.info("New client connected")
options = websocket.recv()
options = json.loads(options)
if self.client_manager is None:
max_clients = options.get('max_clients', 4)
max_connection_time = options.get('max_connection_time', 600)
@@ -785,10 +789,7 @@ class ServeClientFasterWhisper(ServeClientBase):
"large-v3-turbo", "turbo"
]
if not os.path.exists(model):
self.model_size_or_path = self.check_valid_model(model)
else:
self.model_size_or_path = model
self.model_size_or_path = model
self.language = "en" if self.model_size_or_path.endswith("en") else language
self.task = task
self.initial_prompt = initial_prompt
@@ -806,15 +807,25 @@ class ServeClientFasterWhisper(ServeClientBase):
if self.model_size_or_path is None:
return
logging.info(f"Using Device={device} with precision {self.compute_type}")
if single_model:
if ServeClientFasterWhisper.SINGLE_MODEL is None:
self.create_model(device)
ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
try:
if single_model:
if ServeClientFasterWhisper.SINGLE_MODEL is None:
self.create_model(device)
ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
else:
self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
else:
self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
else:
self.create_model(device)
self.create_model(device)
except Exception as e:
logging.error(f"Failed to load model: {e}")
self.websocket.send(json.dumps({
"uid": self.client_uid,
"status": "ERROR",
"message": f"Failed to load model: {str(self.model_size_or_path)}"
}))
self.websocket.close()
return
self.use_vad = use_vad