130 lines
6.5 KiB
Python
130 lines
6.5 KiB
Python
import argparse
|
|
import os
|
|
import soundfile as sf
|
|
import random
|
|
import numpy as np
|
|
import torch
|
|
|
|
from vllm.sampling_params import SamplingParams
|
|
|
|
import os as _os_env_toggle
|
|
_os_env_toggle.environ["VLLM_USE_V1"] = "1"
|
|
|
|
from vllm_omni.entrypoints.omni_llm import OmniLLM
|
|
from utils import make_omni_prompt
|
|
|
|
|
|
SEED = 42
|
|
# Set all random seeds
|
|
random.seed(SEED)
|
|
np.random.seed(SEED)
|
|
torch.manual_seed(SEED)
|
|
torch.cuda.manual_seed(SEED)
|
|
torch.cuda.manual_seed_all(SEED)
|
|
|
|
# Make PyTorch deterministic
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
|
|
# Set environment variables for deterministic behavior
|
|
os.environ["PYTHONHASHSEED"] = str(SEED)
|
|
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--model', required=True, help='Path to merged model directory (will be created if downloading).')
|
|
parser.add_argument('--thinker-model', type=str, default=None)
|
|
parser.add_argument('--talker-model', type=str, default=None)
|
|
parser.add_argument('--code2wav-model', type=str, default=None)
|
|
parser.add_argument('--hf-hub-id', default='Qwen/Qwen2.5-Omni-7B', help='Hugging Face repo id to download if needed.')
|
|
parser.add_argument('--hf-revision', default=None, help='Optional HF revision (branch/tag/commit).')
|
|
parser.add_argument('--prompts', required=True, nargs='+', help='Input text prompts.')
|
|
parser.add_argument('--voice-type', default='default', help='Voice type, e.g., m02, f030, default.')
|
|
parser.add_argument('--code2wav-dir', default=None, help='Path to code2wav folder (contains spk_dict.pt).')
|
|
parser.add_argument('--dit-ckpt', default=None, help='Path to DiT checkpoint file (e.g., dit.pt).')
|
|
parser.add_argument('--bigvgan-ckpt', default=None, help='Path to BigVGAN checkpoint file.')
|
|
parser.add_argument('--dtype', default='bfloat16', choices=['float16', 'bfloat16', 'float32'])
|
|
parser.add_argument('--max-model-len', type=int, default=32768)
|
|
|
|
parser.add_argument("--thinker-only", action="store_true")
|
|
parser.add_argument("--text-only", action="store_true")
|
|
parser.add_argument("--do-wave", action="store_true")
|
|
parser.add_argument('--prompt_type',
|
|
choices=[
|
|
'text', 'audio', 'audio-long', 'audio-long-chunks',
|
|
'audio-long-expand-chunks', 'image', 'video',
|
|
'video-frames', 'audio-in-video', 'audio-in-video-v2',
|
|
"audio-multi-round", "badcase-vl", "badcase-text",
|
|
"badcase-image-early-stop", "badcase-two-audios",
|
|
"badcase-two-videos", "badcase-multi-round",
|
|
"badcase-voice-type", "badcase-voice-type-v2",
|
|
"badcase-audio-tower-1", "badcase-audio-only"
|
|
],
|
|
default='text')
|
|
parser.add_argument('--use-torchvision', action='store_true')
|
|
parser.add_argument('--tokenize', action='store_true')
|
|
parser.add_argument('--output-wav', default="output.wav", help='Output wav file path.')
|
|
parser.add_argument('--thinker-hidden-states-dir', default="thinker_hidden_states", help='Path to thinker hidden states directory.')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
model_name = args.model
|
|
omni_llm = OmniLLM(model=model_name)
|
|
thinker_sampling_params = SamplingParams(
|
|
temperature=0.0, # Deterministic - no randomness
|
|
top_p=1.0, # Disable nucleus sampling
|
|
top_k=-1, # Disable top-k sampling
|
|
max_tokens=2048,
|
|
seed=SEED, # Fixed seed for sampling
|
|
detokenize=True,
|
|
repetition_penalty=1.1,
|
|
)
|
|
talker_sampling_params = SamplingParams(
|
|
temperature=0.0, # Deterministic - no randomness
|
|
top_p=1.0, # Disable nucleus sampling
|
|
top_k=-1, # Disable top-k sampling
|
|
max_tokens=2048,
|
|
seed=SEED, # Fixed seed for sampling
|
|
detokenize=True,
|
|
repetition_penalty=1.1,
|
|
stop_token_ids=[8294]
|
|
)
|
|
code2wav_sampling_params = SamplingParams(
|
|
temperature=0.0, # Deterministic - no randomness
|
|
top_p=1.0, # Disable nucleus sampling
|
|
top_k=-1, # Disable top-k sampling
|
|
max_tokens=2048,
|
|
seed=SEED, # Fixed seed for sampling
|
|
detokenize=True,
|
|
repetition_penalty=1.1,
|
|
)
|
|
|
|
sampling_params_list = [thinker_sampling_params,
|
|
talker_sampling_params,
|
|
code2wav_sampling_params]
|
|
|
|
prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
|
|
omni_outputs = omni_llm.generate(prompt, sampling_params_list)
|
|
|
|
os.makedirs(args.output_wav, exist_ok=True)
|
|
for stage_outputs in omni_outputs:
|
|
if stage_outputs.final_output_type == "text":
|
|
for output in stage_outputs.request_output:
|
|
request_id = output.request_id
|
|
text_output = output.outputs[0].text
|
|
print(f"Request ID: {request_id}, Text Output: {text_output}")
|
|
elif stage_outputs.final_output_type == "audio":
|
|
for output in stage_outputs.request_output:
|
|
request_id = output.request_id
|
|
audio_tensor = output.multimodal_output["audio"]
|
|
output_wav = os.path.join(args.output_wav, f"output_{output.request_id}.wav")
|
|
sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
|
|
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |