Add VoiceClonePrompt.save()/load() for cross-session voice reuse

[HenryVarro666]

This adds VoiceClonePrompt.save(path) and VoiceClonePrompt.load(path) so a cloned voice can be persisted and reused across sessions, plus a short README example. Refs #54, #169, #44.

Motivation

In #169 you pointed out that create_voice_clone_prompt() can be reused across generate() calls — that works within one session, but the prompt cannot be persisted, so every new session pays the full cost again: audio loading, silence trimming, tokenization, and (when ref_text is omitted) loading Whisper and running ASR. For the batch use case in #169 (1k+ files over multiple runs) this is significant, and because auto-transcription can vary slightly between runs, recomputing also means the prompt is not reproducible. #54's ask for saved voices needs an on-disk format as its first building block.

Why an official format instead of "just torch.save the object"

With torch >= 2.6 (weights_only=True became the torch.load default; this repo pins torch 2.8), pickling the dataclass directly produces a file that fails to load by default:

torch.save(prompt, "naive.pt")
torch.load("naive.pt")
# _pickle.UnpicklingError: Weights only load failed. ...
# Re-running `torch.load` with `weights_only` set to `False` will likely succeed,
# but it can result in arbitrary code execution. ...

So the DIY route pushes users toward weights_only=False, which is exactly the footgun the torch default is trying to prevent. This PR instead saves a version-tagged plain dict (tensor moved to CPU), which loads cleanly under weights_only=True and is portable across devices (CUDA/MPS/CPU).

Design notes

• Pure addition: no existing code path changes, no new dependencies, no change to generate() behavior.
• Tokens are stored on CPU; on load, nothing needs to be moved manually — _prepare_inference_inputs already does .to(self.device) when the prompt is consumed.
• format_version field allows the format to evolve without breaking old files.
• VoiceClonePrompt is now exported from omnivoice so the README example works with a plain import.
• The voice-library manager in feat: streaming TTS generation and persistent voice library #166 covers a broader scope (named library, demo UI); this PR only adds the minimal serialization primitive on the existing dataclass, which such a library could build on.

Verification (Apple Silicon, MPS, torch 2.8.0)

Session 1 — create, save, round-trip check, generate:

model = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="mps", dtype=torch.float16)
p = model.create_voice_clone_prompt(ref_audio="ref.wav")   # 228s incl. Whisper download + ASR
p.save("my_voice.pt")
q = VoiceClonePrompt.load("my_voice.pt")
assert torch.equal(p.ref_audio_tokens.cpu(), q.ref_audio_tokens)  # bitwise identical
assert p.ref_text == q.ref_text and p.ref_rms == q.ref_rms
audio = model.generate(text="This voice was saved and loaded from disk.", voice_clone_prompt=q)  # OK

Session 2 — fresh process:

model = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="mps", dtype=torch.float16)
p = VoiceClonePrompt.load("my_voice.pt")                   # 1.4 ms
audio = model.generate(text="Loaded in a completely new session.", voice_clone_prompt=p)  # OK
assert model._asr_pipe is None                              # Whisper never loaded

The saved prompt is byte-identical across sessions, and the second session never loads the ASR model at all.

pre-commit run --all-files passes.

[gemini-code-assist]

Code Review

This pull request introduces serialization support for VoiceClonePrompt by adding save and load methods, allowing cloned voices to be reused across sessions. It also updates the documentation and exports VoiceClonePrompt in the package root. The review feedback suggests improving the robustness and usability of these methods by supporting os.PathLike and torch.device types, and adding defensive validation checks on the loaded data structure.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

To improve usability and compatibility with standard Python path handling (such as pathlib.Path), consider updating the path parameter type hint to accept Union[str, os.PathLike] instead of just str.

Suggested change

	def save(self, path: str) -> None:
	def save(self, path: Union[str, os.PathLike]) -> None:

[gemini-code-assist]

To make the load method more robust and flexible, we should:

1. Update the path parameter type hint to accept Union[str, os.PathLike] to support pathlib.Path objects.
2. Update map_location to accept Optional[Union[str, torch.device]] to allow loading directly to a specific PyTorch device object.
3. Add defensive checks to ensure the loaded data is a dictionary and contains all the required keys, preventing unexpected AttributeError or unhandled KeyError when loading invalid or corrupted files.

    @classmethod
    def load(
        cls,
        path: Union[str, os.PathLike],
        map_location: Optional[Union[str, torch.device]] = "cpu",
    ) -> "VoiceClonePrompt":
        """Load a prompt saved with :meth:`save`.

        The returned prompt can be passed directly to
        :meth:`OmniVoice.generate`; the audio tokens are moved to the model
        device automatically during generation, so no manual ``.to(device)``
        is needed.

        Args:
            path: File path previously written by :meth:`save`.
            map_location: Device to load the audio tokens onto.
        Returns:
            The restored :class:`VoiceClonePrompt`.
        """
        data = torch.load(path, map_location=map_location, weights_only=True)
        if not isinstance(data, dict):
            raise ValueError(
                f"Invalid VoiceClonePrompt file: expected a dictionary, got {type(data).__name__}"
            )
        version = data.get("format_version")
        if version != _VOICE_CLONE_PROMPT_FORMAT_VERSION:
            raise ValueError(f"Unsupported VoiceClonePrompt format version: {version}")
        
        for key in ("ref_audio_tokens", "ref_text", "ref_rms"):
            if key not in data:
                raise ValueError(f"Invalid VoiceClonePrompt file: missing required key '{key}'")

        return cls(
            ref_audio_tokens=data["ref_audio_tokens"],
            ref_text=data["ref_text"],
            ref_rms=data["ref_rms"],
        )

[zhu-han]

Thanks! Merged