Function process_init

omnivoice/cli/infer_batch.py:207–256

Initializer for each worker process. Loads model (with tokenizers and duration estimator) onto a specific GPU via ``OmniVoice.from_pretrained()``.

(rank_queue, model_checkpoint, warmup=0)

def process_init(rank_queue, model_checkpoint, warmup=0):
    """Initializer for each worker process.

    Loads model (with tokenizers and duration estimator) onto a specific GPU
    via ``OmniVoice.from_pretrained()``.
    """
    global worker_model

    torch.set_num_threads(2)
    torch.set_num_interop_threads(2)

    formatter = (
        "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] "
        "[Worker %(process)d] %(message)s"
    )
    logging.basicConfig(format=formatter, level=logging.INFO, force=True)

    rank = rank_queue.get()
    device_type, device_id = rank
    if device_type == "cpu":
        worker_device = "cpu"
    elif device_type == "mps":
        worker_device = "mps"
    else:
        worker_device = f"cuda:{device_id}"

    logging.info(f"Initializing worker on device: {worker_device}")

    worker_model = OmniVoice.from_pretrained(
        model_checkpoint,
        device_map=worker_device,
        dtype=torch.float16,
    )

    if warmup > 0:
        logging.info(f"Running {warmup} warmup iterations on {worker_device}")
        dummy_ref_audio = (
            torch.randn(1, SAMPLING_RATE),
            SAMPLING_RATE,
        )  # 1s dummy audio
        for i in range(warmup):
            worker_model.generate(
                text=["hello"],
                language=["en"],
                ref_audio=[dummy_ref_audio],
                ref_text=["hello"],
            )
        logging.info(f"Warmup complete on {worker_device}")

    logging.info(f"Worker on {worker_device} initialized successfully.")


def estimate_sample_total_duration(