From 0897414705b132ea80e15a75978b487bb78cb8e7 Mon Sep 17 00:00:00 2001
From: Chao Cao <1504517223@qq.com>
Date: Mon, 13 Jul 2026 10:44:51 -0400
Subject: [PATCH] Add VoiceClonePrompt.save()/load() for cross-session voice
 reuse

---
 README.md                     | 18 +++++++++++++
 omnivoice/__init__.py         |  8 +++++-
 omnivoice/models/omnivoice.py | 51 ++++++++++++++++++++++++++++++++++-
 3 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index b6f54e7e..2adcb008 100644
--- a/README.md
+++ b/README.md
@@ -168,6 +168,24 @@ audio = model.generate(
 sf.write("out.wav", audio[0], 24000)
 ```
 
+#### Reusing a cloned voice across sessions
+
+Encode the reference audio once, save the resulting prompt, and skip the
+audio loading / auto-transcription steps in later sessions:
+
+```python
+prompt = model.create_voice_clone_prompt(
+    ref_audio="ref.wav", ref_text="Transcription of the reference audio."
+)
+prompt.save("my_voice.pt")
+
+# Later, in a new session:
+from omnivoice import VoiceClonePrompt
+
+prompt = VoiceClonePrompt.load("my_voice.pt")
+audio = model.generate(text="Hello again!", voice_clone_prompt=prompt)
+```
+
 > **Tips**
 >
 > - Use a 3–10 seconds reference audio clip. Longer audio slows down inference and may degrade cloning quality.
diff --git a/omnivoice/__init__.py b/omnivoice/__init__.py
index fb755bf7..c26337fc 100644
--- a/omnivoice/__init__.py
+++ b/omnivoice/__init__.py
@@ -23,6 +23,12 @@
     OmniVoice,
     OmniVoiceConfig,
     OmniVoiceGenerationConfig,
+    VoiceClonePrompt,
 )
 
-__all__ = ["OmniVoice", "OmniVoiceConfig", "OmniVoiceGenerationConfig"]
+__all__ = [
+    "OmniVoice",
+    "OmniVoiceConfig",
+    "OmniVoiceGenerationConfig",
+    "VoiceClonePrompt",
+]
diff --git a/omnivoice/models/omnivoice.py b/omnivoice/models/omnivoice.py
index 51dd8bd8..ff68acae 100644
--- a/omnivoice/models/omnivoice.py
+++ b/omnivoice/models/omnivoice.py
@@ -113,12 +113,60 @@ def _autocast_flex_attention(module, query, key, value, *args, **kwargs):
 # ---------------------------------------------------------------------------
 
 
+_VOICE_CLONE_PROMPT_FORMAT_VERSION = 1
+
+
 @dataclass
 class VoiceClonePrompt:
     ref_audio_tokens: torch.Tensor  # (C, T)
     ref_text: str
     ref_rms: float
 
+    def save(self, path: str) -> None:
+        """Save this prompt to ``path`` for reuse in a later session.
+
+        The file stores a plain dict with the audio tokens moved to CPU, so
+        it can be loaded with ``torch.load(weights_only=True)`` (the default
+        since torch 2.6) and is portable across devices.
+
+        Args:
+            path: Destination file path (e.g. ``"my_voice.pt"``).
+        """
+        torch.save(
+            {
+                "format_version": _VOICE_CLONE_PROMPT_FORMAT_VERSION,
+                "ref_audio_tokens": self.ref_audio_tokens.detach().cpu(),
+                "ref_text": self.ref_text,
+                "ref_rms": float(self.ref_rms),
+            },
+            path,
+        )
+
+    @classmethod
+    def load(cls, path: str, map_location: str = "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)
+        version = data.get("format_version")
+        if version != _VOICE_CLONE_PROMPT_FORMAT_VERSION:
+            raise ValueError(f"Unsupported VoiceClonePrompt format version: {version}")
+        return cls(
+            ref_audio_tokens=data["ref_audio_tokens"],
+            ref_text=data["ref_text"],
+            ref_rms=data["ref_rms"],
+        )
+
 
 @dataclass
 class OmniVoiceGenerationConfig:
@@ -555,7 +603,8 @@ def generate(
             ref_text: Optional reference text for voice cloning mode.
             ref_audio: Optional reference audio for voice cloning mode.
                 Can be a file path or a (waveform, sample_rate) tuple.
-            voice_clone_prompt: Reusable prompt from :meth:`create_voice_clone_prompt`.
+            voice_clone_prompt: Reusable prompt from :meth:`create_voice_clone_prompt`
+                or :meth:`VoiceClonePrompt.load`.
                 If provided, it overrides ``ref_text`` and ``ref_audio``.
             instruct: Style instruction for voice design mode.
             duration: Fixed output duration in seconds. If a single float,