Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
8 changes: 7 additions & 1 deletion omnivoice/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,12 @@
OmniVoice,
OmniVoiceConfig,
OmniVoiceGenerationConfig,
VoiceClonePrompt,
)

__all__ = ["OmniVoice", "OmniVoiceConfig", "OmniVoiceGenerationConfig"]
__all__ = [
"OmniVoice",
"OmniVoiceConfig",
"OmniVoiceGenerationConfig",
"VoiceClonePrompt",
]
51 changes: 50 additions & 1 deletion omnivoice/models/omnivoice.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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:

"""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"],
)
Comment on lines +146 to +168

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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"],
        )



@dataclass
class OmniVoiceGenerationConfig:
Expand Down Expand Up @@ -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,
Expand Down