Inital Commit

a42ddbb over 1 year ago

4.96 kB

# model_manager.py

import torch
from TTS.api import TTS
from voice.clearvoice import ClearVoice
import os

os.environ["COQUI_TOS_AGREED"] = "1"

_model_cache = {}


def load_model():
    global _model_cache
    """Load and cache all necessary models."""
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if "xtts_v2" not in _model_cache:
        _model_cache["xtts_v2"] = TTS(
            "tts_models/multilingual/multi-dataset/xtts_v2"
        ).to(device)

    if "FRCRN_SE_16K" not in _model_cache:
        _model_cache["FRCRN_SE_16K"] = ClearVoice(
            task="speech_enhancement", model_names=["FRCRN_SE_16K"]
        )

    if "MossFormer2_SE_48K" not in _model_cache:
        _model_cache["MossFormer2_SE_48K"] = ClearVoice(
            task="speech_enhancement", model_names=["MossFormer2_SE_48K"]
        )

    if "MossFormer2_SS_16K" not in _model_cache:
        _model_cache["MossFormer2_SS_16K"] = ClearVoice(
            task="speech_separation", model_names=["MossFormer2_SS_16K"]
        )

    if "AV_MossFormer2_TSE_16K" not in _model_cache:
        _model_cache["AV_MossFormer2_TSE_16K"] = ClearVoice(
            task="target_speaker_extraction", model_names=["AV_MossFormer2_TSE_16K"]
        )

    if "MossFormer2_SR_48K" not in _model_cache:
        _model_cache["MossFormer2_SR_48K"] = ClearVoice(
            task="speech_super_resolution", model_names=["MossFormer2_SR_48K"]
        )

    return _model_cache


def unload_model():
    """Clear cached models and free up memory."""
    _model_cache.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def predict(task: str, **kwargs):
    """
    Central prediction dispatcher based on task type.
    Args:
        task (str): Task type key.
        kwargs: Task-specific arguments.
    Returns:
        Result of the model inference.
    """
    if not _model_cache:
        load_model()

    if task == "clone_voice":
        model = _model_cache["xtts_v2"]
        return model.tts_to_file(
            text=kwargs["text"],
            speaker_wav=kwargs["speaker_wav"],
            language=kwargs.get("language", "en"),
            file_path=kwargs["file_path"],
        )

    elif task == "speech_enhancement_16k":
        model = _model_cache["FRCRN_SE_16K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
        )

    elif task == "speech_enhancement_48k":
        model = _model_cache["MossFormer2_SE_48K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
            output_path=kwargs.get("output_path"),
        )

    elif task == "speech_separation":
        model = _model_cache["MossFormer2_SS_16K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
        )

    elif task == "target_speaker_extraction":
        model = _model_cache["AV_MossFormer2_TSE_16K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=True,
            output_path=kwargs["output_path"],
        )

    elif task == "speech_super_resolution":
        model = _model_cache["MossFormer2_SR_48K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
        )

    else:
        raise ValueError(f"Unknown task type: {task}")


# from model_manager import predict

# # Assume you have speaker and output audio in bytes
# with open("speaker.wav", "rb") as f:
#     speaker_bytes = f.read()

# text_to_speak = "Hello, this is a cloned voice."
# output_path = "cloned_voice.wav"

# predict(
#     task="clone_voice",
#     text=text_to_speak,
#     speaker_wav=speaker_bytes,
#     language="en",
#     file_path=output_path
# )


# from model_manager import predict

# with open("noisy_16k.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_enhancement_16k",
#     input_audio=audio_bytes
# )


# from model_manager import predict

# with open("noisy_48k.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_enhancement_48k",
#     input_audio=audio_bytes,
#     output_path="enhanced_48k.wav"
# )


# from model_manager import predict

# with open("overlapping_speakers.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_separation",
#     input_audio=audio_bytes
# )


# from model_manager import predict

# with open("mixture_with_target.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="target_speaker_extraction",
#     input_audio=audio_bytes,
#     output_path="target_extracted.wav"
# )


# from model_manager import predict

# with open("low_quality_48k.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_super_resolution",
#     input_audio=audio_bytes
# )

	# model_manager.py

	import torch
	from TTS.api import TTS
	from voice.clearvoice import ClearVoice
	import os

	os.environ["COQUI_TOS_AGREED"] = "1"

	_model_cache = {}


	def load_model():
	global _model_cache
	"""Load and cache all necessary models."""
	device = "cuda" if torch.cuda.is_available() else "cpu"

	if "xtts_v2" not in _model_cache:
	_model_cache["xtts_v2"] = TTS(
	"tts_models/multilingual/multi-dataset/xtts_v2"
	).to(device)

	if "FRCRN_SE_16K" not in _model_cache:
	_model_cache["FRCRN_SE_16K"] = ClearVoice(
	task="speech_enhancement", model_names=["FRCRN_SE_16K"]
	)

	if "MossFormer2_SE_48K" not in _model_cache:
	_model_cache["MossFormer2_SE_48K"] = ClearVoice(
	task="speech_enhancement", model_names=["MossFormer2_SE_48K"]
	)

	if "MossFormer2_SS_16K" not in _model_cache:
	_model_cache["MossFormer2_SS_16K"] = ClearVoice(
	task="speech_separation", model_names=["MossFormer2_SS_16K"]
	)

	if "AV_MossFormer2_TSE_16K" not in _model_cache:
	_model_cache["AV_MossFormer2_TSE_16K"] = ClearVoice(
	task="target_speaker_extraction", model_names=["AV_MossFormer2_TSE_16K"]
	)

	if "MossFormer2_SR_48K" not in _model_cache:
	_model_cache["MossFormer2_SR_48K"] = ClearVoice(
	task="speech_super_resolution", model_names=["MossFormer2_SR_48K"]
	)

	return _model_cache


	def unload_model():
	"""Clear cached models and free up memory."""
	_model_cache.clear()
	if torch.cuda.is_available():
	torch.cuda.empty_cache()


	def predict(task: str, **kwargs):
	"""
	Central prediction dispatcher based on task type.

	Args:
	task (str): Task type key.
	kwargs: Task-specific arguments.

	Returns:
	Result of the model inference.
	"""
	if not _model_cache:
	load_model()

	if task == "clone_voice":
	model = _model_cache["xtts_v2"]
	return model.tts_to_file(
	text=kwargs["text"],
	speaker_wav=kwargs["speaker_wav"],
	language=kwargs.get("language", "en"),
	file_path=kwargs["file_path"],
	)

	elif task == "speech_enhancement_16k":
	model = _model_cache["FRCRN_SE_16K"]
	return model(
	input_path=kwargs["input_path"],
	online_write=kwargs.get("online_write", False),
	)

	elif task == "speech_enhancement_48k":
	model = _model_cache["MossFormer2_SE_48K"]
	return model(
	input_path=kwargs["input_path"],
	online_write=kwargs.get("online_write", False),
	output_path=kwargs.get("output_path"),
	)

	elif task == "speech_separation":
	model = _model_cache["MossFormer2_SS_16K"]
	return model(
	input_path=kwargs["input_path"],
	online_write=kwargs.get("online_write", False),
	)

	elif task == "target_speaker_extraction":
	model = _model_cache["AV_MossFormer2_TSE_16K"]
	return model(
	input_path=kwargs["input_path"],
	online_write=True,
	output_path=kwargs["output_path"],
	)

	elif task == "speech_super_resolution":
	model = _model_cache["MossFormer2_SR_48K"]
	return model(
	input_path=kwargs["input_path"],
	online_write=kwargs.get("online_write", False),
	)

	else:
	raise ValueError(f"Unknown task type: {task}")


	# from model_manager import predict

	# # Assume you have speaker and output audio in bytes
	# with open("speaker.wav", "rb") as f:
	# speaker_bytes = f.read()

	# text_to_speak = "Hello, this is a cloned voice."
	# output_path = "cloned_voice.wav"

	# predict(
	# task="clone_voice",
	# text=text_to_speak,
	# speaker_wav=speaker_bytes,
	# language="en",
	# file_path=output_path
	# )


	# from model_manager import predict

	# with open("noisy_16k.wav", "rb") as f:
	# audio_bytes = f.read()

	# predict(
	# task="speech_enhancement_16k",
	# input_audio=audio_bytes
	# )


	# from model_manager import predict

	# with open("noisy_48k.wav", "rb") as f:
	# audio_bytes = f.read()

	# predict(
	# task="speech_enhancement_48k",
	# input_audio=audio_bytes,
	# output_path="enhanced_48k.wav"
	# )


	# from model_manager import predict

	# with open("overlapping_speakers.wav", "rb") as f:
	# audio_bytes = f.read()

	# predict(
	# task="speech_separation",
	# input_audio=audio_bytes
	# )


	# from model_manager import predict

	# with open("mixture_with_target.wav", "rb") as f:
	# audio_bytes = f.read()

	# predict(
	# task="target_speaker_extraction",
	# input_audio=audio_bytes,
	# output_path="target_extracted.wav"
	# )


	# from model_manager import predict

	# with open("low_quality_48k.wav", "rb") as f:
	# audio_bytes = f.read()

	# predict(
	# task="speech_super_resolution",
	# input_audio=audio_bytes
	# )