diff --git a/README.md b/README.md
index 063f330..f5e63ff 100644
--- a/README.md
+++ b/README.md
@@ -1,58 +1,23 @@
-
-
-

+# ComfyUI OpenVoice (Unofficial Support)
-[Paper](https://arxiv.org/abs/2312.01479) |
-[Website](https://research.myshell.ai/open-voice)
+Unofficial [OpenVoice](https://github.com/myshell-ai/OpenVoice) support for [ComfyUI](https://github.com/comfyanonymous/ComfyUI).
-
+## Important Updates
+- **2024/04/12:** Added ComfyUI nodes and workflow examples
-## Introduction
-As we detailed in our [paper](https://arxiv.org/abs/2312.01479) and [website](https://research.myshell.ai/open-voice), the advantages of OpenVoice are three-fold:
+## Basic Workflow
+This [workflow](examples/workflow.json) shows the basic usages on (1) text-to-speech (TTS) and (2) speech-to-speech (STS) using reference voice.
-**1. Accurate Tone Color Cloning.**
-OpenVoice can accurately clone the reference tone color and generate speech in multiple languages and accents.
+ 
-**2. Flexible Voice Style Control.**
-OpenVoice enables granular control over voice styles, such as emotion and accent, as well as other style parameters including rhythm, pauses, and intonation.
+## Installation
+- (WIP) Install from ComfyUI Manager (search for `openvoice`, make sure `ffmpeg` is installed)
-**3. Zero-shot Cross-lingual Voice Cloning.**
-Neither of the language of the generated speech nor the language of the reference speech needs to be presented in the massive-speaker multi-lingual training dataset.
-
-[Video](https://github.com/myshell-ai/OpenVoice/assets/40556743/3cba936f-82bf-476c-9e52-09f0f417bb2f)
-
-OpenVoice has been powering the instant voice cloning capability of [myshell.ai](https://app.myshell.ai/explore) since May 2023. Until Nov 2023, the voice cloning model has been used tens of millions of times by users worldwide, and witnessed the explosive user growth on the platform.
-
-## Main Contributors
-
-- [Zengyi Qin](https://www.qinzy.tech) at MIT and MyShell
-- [Wenliang Zhao](https://wl-zhao.github.io) at Tsinghua University
-- [Xumin Yu](https://yuxumin.github.io) at Tsinghua University
-- [Ethan Sun](https://twitter.com/ethan_myshell) at MyShell
-
-## How to Use
-Please see [usage](docs/USAGE.md) for detailed instructions.
-
-## Common Issues
-
-Please see [QA](docs/QA.md) for common questions and answers. We will regularly update the question and answer list.
-
-## Join Our Community
-
-Join our [Discord community](https://discord.gg/myshell) and select the `Developer` role upon joining to gain exclusive access to our developer-only channel! Don't miss out on valuable discussions and collaboration opportunities.
-
-## Citation
+- Download or git clone this repository into the ComfyUI/custom_nodes/ directory and run:
```
-@article{qin2023openvoice,
- title={OpenVoice: Versatile Instant Voice Cloning},
- author={Qin, Zengyi and Zhao, Wenliang and Yu, Xumin and Sun, Xin},
- journal={arXiv preprint arXiv:2312.01479},
- year={2023}
-}
+sudo apt install ffmpeg
+pip install -r requirements.txt
```
-## License
-This repository is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License, which prohibits commercial usage. **This will be changed to a license that allows Free Commercial usage in the near future.** Stay tuned. For social responsibility and anti-misuse considerations, **MyShell reserves the ability to detect whether an audio is generated by OpenVoice**, no matter whether the watermark is added or not.
-
-## Acknowledgements
-This implementation is based on several excellent projects, [TTS](https://github.com/coqui-ai/TTS), [VITS](https://github.com/jaywalnut310/vits), and [VITS2](https://github.com/daniilrobnikov/vits2). Thanks for their awesome work!
+## Download Checkpoints
+All the models are downloaded automatically, no worries.
diff --git a/README.md.bak b/README.md.bak
new file mode 100644
index 0000000..063f330
--- /dev/null
+++ b/README.md.bak
@@ -0,0 +1,58 @@
+
+
+

+
+[Paper](https://arxiv.org/abs/2312.01479) |
+[Website](https://research.myshell.ai/open-voice)
+
+
+
+## Introduction
+As we detailed in our [paper](https://arxiv.org/abs/2312.01479) and [website](https://research.myshell.ai/open-voice), the advantages of OpenVoice are three-fold:
+
+**1. Accurate Tone Color Cloning.**
+OpenVoice can accurately clone the reference tone color and generate speech in multiple languages and accents.
+
+**2. Flexible Voice Style Control.**
+OpenVoice enables granular control over voice styles, such as emotion and accent, as well as other style parameters including rhythm, pauses, and intonation.
+
+**3. Zero-shot Cross-lingual Voice Cloning.**
+Neither of the language of the generated speech nor the language of the reference speech needs to be presented in the massive-speaker multi-lingual training dataset.
+
+[Video](https://github.com/myshell-ai/OpenVoice/assets/40556743/3cba936f-82bf-476c-9e52-09f0f417bb2f)
+
+OpenVoice has been powering the instant voice cloning capability of [myshell.ai](https://app.myshell.ai/explore) since May 2023. Until Nov 2023, the voice cloning model has been used tens of millions of times by users worldwide, and witnessed the explosive user growth on the platform.
+
+## Main Contributors
+
+- [Zengyi Qin](https://www.qinzy.tech) at MIT and MyShell
+- [Wenliang Zhao](https://wl-zhao.github.io) at Tsinghua University
+- [Xumin Yu](https://yuxumin.github.io) at Tsinghua University
+- [Ethan Sun](https://twitter.com/ethan_myshell) at MyShell
+
+## How to Use
+Please see [usage](docs/USAGE.md) for detailed instructions.
+
+## Common Issues
+
+Please see [QA](docs/QA.md) for common questions and answers. We will regularly update the question and answer list.
+
+## Join Our Community
+
+Join our [Discord community](https://discord.gg/myshell) and select the `Developer` role upon joining to gain exclusive access to our developer-only channel! Don't miss out on valuable discussions and collaboration opportunities.
+
+## Citation
+```
+@article{qin2023openvoice,
+ title={OpenVoice: Versatile Instant Voice Cloning},
+ author={Qin, Zengyi and Zhao, Wenliang and Yu, Xumin and Sun, Xin},
+ journal={arXiv preprint arXiv:2312.01479},
+ year={2023}
+}
+```
+
+## License
+This repository is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License, which prohibits commercial usage. **This will be changed to a license that allows Free Commercial usage in the near future.** Stay tuned. For social responsibility and anti-misuse considerations, **MyShell reserves the ability to detect whether an audio is generated by OpenVoice**, no matter whether the watermark is added or not.
+
+## Acknowledgements
+This implementation is based on several excellent projects, [TTS](https://github.com/coqui-ai/TTS), [VITS](https://github.com/jaywalnut310/vits), and [VITS2](https://github.com/daniilrobnikov/vits2). Thanks for their awesome work!
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..d721463
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..a65ccd3
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,160 @@
+import os
+import torch
+import random
+import folder_paths
+import soundfile as sf
+
+from .openvoice import se_extractor
+from .openvoice.api import BaseSpeakerTTS, ToneColorConverter
+
+
+class AnyType(str):
+ def __eq__(self, _) -> bool:
+ return True
+
+ def __ne__(self, __value: object) -> bool:
+ return False
+
+any = AnyType("*")
+
+class OpenVoiceTTS:
+ @classmethod
+ def INPUT_TYPES(s):
+ audio_extensions = ["wav", "mp3", "flac"]
+ input_dir = folder_paths.get_input_directory()
+ files = []
+ for f in os.listdir(input_dir):
+ if os.path.isfile(os.path.join(input_dir, f)):
+ file_parts = f.lower().split('.')
+ if len(file_parts) > 1 and (file_parts[-1] in audio_extensions):
+ files.append(f)
+ return {
+ "required": {
+ "text": ("STRING", {"default": '', "multiline": True}),
+ "lang": (["English","Chinese"],),
+ "style": (["default","whispering","cheerful","terrified","angry","sad","friendly"],),
+ "speed": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
+ "ref_voice": (sorted(files),),
+ },
+ }
+
+ CATEGORY = "OpenVoice"
+
+ RETURN_TYPES = (any, "INT",)
+ RETURN_NAMES = ("AUDIO", "SAMPLE_RATE",)
+ FUNCTION = "inference"
+
+ def inference(self, text, lang, style, speed, ref_voice):
+ local_dir = os.path.join(folder_paths.models_dir, 'openovice')
+ if not os.path.exists(local_dir) or not os.path.isdir(local_dir):
+ from huggingface_hub import snapshot_download
+ snapshot_download(repo_id="Alignment-Lab-AI/OpenVoice", local_dir=local_dir, local_dir_use_symlinks=False)
+
+ mark = BaseSpeakerTTS.language_marks.get(lang.lower(), None)
+ assert mark is not None, f"language {lang} is not supported"
+
+ ckpt_base = os.path.join(local_dir, f'checkpoints/base_speakers/{mark}')
+ ckpt_converter = os.path.join(local_dir, 'checkpoints/converter')
+ device="cuda:0" if torch.cuda.is_available() else "cpu"
+
+ base_speaker_tts = BaseSpeakerTTS(f'{ckpt_base}/config.json', device=device)
+ base_speaker_tts.load_ckpt(f'{ckpt_base}/checkpoint.pth')
+
+ tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device)
+ tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')
+
+ style_name = 'default' if style == 'default' else 'style'
+ source_se = torch.load(f'{ckpt_base}/{mark.lower()}_{style_name}_se.pth').to(device)
+ reference_speaker = os.path.join(folder_paths.get_input_directory(), ref_voice)
+
+ temp_dir = folder_paths.get_temp_directory()
+ file_prefix = ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5))
+ target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir=temp_dir, vad=True)
+
+ save_path = f'{temp_dir}/{file_prefix}_output_{mark.lower()}_{style}.wav'
+
+ # Run the base speaker tts
+ src_path = f'{temp_dir}/{file_prefix}_base_{mark.lower()}_{style}.wav'
+ base_speaker_tts.tts(text, src_path, speaker=style, language=lang, speed=speed)
+
+ # Run the tone color converter
+ encode_message = "@MyShell"
+ tone_color_converter.convert(
+ audio_src_path=src_path,
+ src_se=source_se,
+ tgt_se=target_se,
+ output_path=save_path,
+ message=encode_message)
+
+ audio_samples, sample_rate =sf.read(save_path)
+ return (list(audio_samples), sample_rate)
+
+class OpenVoiceSTS:
+ @classmethod
+ def INPUT_TYPES(s):
+ audio_extensions = ["wav", "mp3", "flac"]
+ input_dir = folder_paths.get_input_directory()
+ files = []
+ for f in os.listdir(input_dir):
+ if os.path.isfile(os.path.join(input_dir, f)):
+ file_parts = f.lower().split('.')
+ if len(file_parts) > 1 and (file_parts[-1] in audio_extensions):
+ files.append(f)
+ return {
+ "required": {
+ "src_voice": (sorted(files),),
+ "ref_voice": (sorted(files),),
+ },
+ }
+
+ CATEGORY = "OpenVoice"
+
+ RETURN_TYPES = (any, "INT",)
+ RETURN_NAMES = ("AUDIO", "SAMPLE_RATE",)
+ FUNCTION = "inference"
+
+ def inference(self, src_voice, ref_voice):
+ local_dir = os.path.join(folder_paths.models_dir, 'openovice')
+ if not os.path.exists(local_dir) or not os.path.isdir(local_dir):
+ from huggingface_hub import snapshot_download
+ snapshot_download(repo_id="Alignment-Lab-AI/OpenVoice", local_dir=local_dir, local_dir_use_symlinks=False)
+
+ ckpt_converter = os.path.join(local_dir, 'checkpoints/converter')
+ device="cuda:0" if torch.cuda.is_available() else "cpu"
+
+ tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device)
+ tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')
+
+ temp_dir = folder_paths.get_temp_directory()
+ file_prefix = ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5))
+
+ source_speaker = os.path.join(folder_paths.get_input_directory(), src_voice)
+ source_se, audio_name = se_extractor.get_se(source_speaker, tone_color_converter, target_dir=temp_dir, vad=True)
+
+ reference_speaker = os.path.join(folder_paths.get_input_directory(), ref_voice)
+ target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir=temp_dir, vad=True)
+
+ save_path = f'{temp_dir}/{file_prefix}_output_crosslingual.wav'
+
+ # Run the tone color converter
+ encode_message = "@MyShell"
+ tone_color_converter.convert(
+ audio_src_path=source_speaker,
+ src_se=source_se,
+ tgt_se=target_se,
+ output_path=save_path,
+ message=encode_message)
+
+ audio_samples, sample_rate =sf.read(save_path)
+ return (list(audio_samples), sample_rate)
+
+
+NODE_CLASS_MAPPINGS = {
+ "D_OpenVoice_TTS": OpenVoiceTTS,
+ "D_OpenVoice_STS": OpenVoiceSTS,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "D_OpenVoice_TTS": "Open Voice TTS",
+ "D_OpenVoice_STS": "Open Voice STS",
+}
\ No newline at end of file
diff --git a/openvoice/api.py b/openvoice/api.py
index 5ee45f2..4f86ce8 100644
--- a/openvoice/api.py
+++ b/openvoice/api.py
@@ -2,13 +2,13 @@ import torch
import numpy as np
import re
import soundfile
-from openvoice import utils
-from openvoice import commons
+from . import utils
+from . import commons
import os
import librosa
-from openvoice.text import text_to_sequence
-from openvoice.mel_processing import spectrogram_torch
-from openvoice.models import SynthesizerTrn
+from .text import text_to_sequence
+from .mel_processing import spectrogram_torch
+from .models import SynthesizerTrn
class OpenVoiceBaseClass(object):
diff --git a/openvoice/attentions.py b/openvoice/attentions.py
index 5c410f0..3551157 100644
--- a/openvoice/attentions.py
+++ b/openvoice/attentions.py
@@ -3,7 +3,7 @@ import torch
from torch import nn
from torch.nn import functional as F
-from openvoice import commons
+from . import commons
import logging
logger = logging.getLogger(__name__)
diff --git a/openvoice/models.py b/openvoice/models.py
index b1030fd..1ae3165 100644
--- a/openvoice/models.py
+++ b/openvoice/models.py
@@ -3,14 +3,14 @@ import torch
from torch import nn
from torch.nn import functional as F
-from openvoice import commons
-from openvoice import modules
-from openvoice import attentions
+from . import commons
+from . import modules
+from . import attentions
from torch.nn import Conv1d, ConvTranspose1d, Conv2d
from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-from openvoice.commons import init_weights, get_padding
+from .commons import init_weights, get_padding
class TextEncoder(nn.Module):
diff --git a/openvoice/modules.py b/openvoice/modules.py
index d659a32..5af5d11 100644
--- a/openvoice/modules.py
+++ b/openvoice/modules.py
@@ -6,10 +6,10 @@ from torch.nn import functional as F
from torch.nn import Conv1d
from torch.nn.utils import weight_norm, remove_weight_norm
-from openvoice import commons
-from openvoice.commons import init_weights, get_padding
-from openvoice.transforms import piecewise_rational_quadratic_transform
-from openvoice.attentions import Encoder
+from . import commons
+from .commons import init_weights, get_padding
+from .transforms import piecewise_rational_quadratic_transform
+from .attentions import Encoder
LRELU_SLOPE = 0.1
diff --git a/openvoice/openvoice_app.py b/openvoice/openvoice_app.py
index 15b0b43..744a767 100644
--- a/openvoice/openvoice_app.py
+++ b/openvoice/openvoice_app.py
@@ -4,8 +4,8 @@ import argparse
import gradio as gr
from zipfile import ZipFile
import langid
-from openvoice import se_extractor
-from openvoice.api import BaseSpeakerTTS, ToneColorConverter
+from . import se_extractor
+from .api import BaseSpeakerTTS, ToneColorConverter
parser = argparse.ArgumentParser()
parser.add_argument("--share", action='store_true', default=False, help="make link public")
diff --git a/openvoice/text/__init__.py b/openvoice/text/__init__.py
index 6494bcb..9e23777 100644
--- a/openvoice/text/__init__.py
+++ b/openvoice/text/__init__.py
@@ -1,6 +1,6 @@
""" from https://github.com/keithito/tacotron """
-from openvoice.text import cleaners
-from openvoice.text.symbols import symbols
+from . import cleaners
+from .symbols import symbols
# Mappings from symbol to numeric ID and vice versa:
@@ -43,7 +43,7 @@ def cleaned_text_to_sequence(cleaned_text, symbols):
-from openvoice.text.symbols import language_tone_start_map
+from .symbols import language_tone_start_map
def cleaned_text_to_sequence_vits2(cleaned_text, tones, language, symbols, languages):
"""Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
Args:
diff --git a/openvoice/text/cleaners.py b/openvoice/text/cleaners.py
index 16dd168..d58aeab 100644
--- a/openvoice/text/cleaners.py
+++ b/openvoice/text/cleaners.py
@@ -1,6 +1,6 @@
import re
-from openvoice.text.english import english_to_lazy_ipa, english_to_ipa2, english_to_lazy_ipa2
-from openvoice.text.mandarin import number_to_chinese, chinese_to_bopomofo, latin_to_bopomofo, chinese_to_romaji, chinese_to_lazy_ipa, chinese_to_ipa, chinese_to_ipa2
+from .english import english_to_lazy_ipa, english_to_ipa2, english_to_lazy_ipa2
+from .mandarin import number_to_chinese, chinese_to_bopomofo, latin_to_bopomofo, chinese_to_romaji, chinese_to_lazy_ipa, chinese_to_ipa, chinese_to_ipa2
def cjke_cleaners2(text):
text = re.sub(r'\[ZH\](.*?)\[ZH\]',
diff --git a/requirements.txt b/requirements.txt
index 8ddba70..e213e82 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,16 +1,14 @@
-librosa==0.9.1
-faster-whisper==0.9.0
-pydub==0.25.1
-wavmark==0.0.3
-numpy==1.22.0
-eng_to_ipa==0.0.2
-inflect==7.0.0
-unidecode==1.3.7
-whisper-timestamped==1.14.2
-openai
-python-dotenv
-pypinyin==0.50.0
-cn2an==0.5.22
-jieba==0.42.1
-gradio==3.48.0
-langid==1.1.6
+torch
+librosa
+pydub
+faster-whisper
+whisper-timestamped
+inflect
+unidecode
+eng_to_ipa
+pypinyin
+jieba
+cn2an
+wavmark
+huggingface_hub
+soundfile
\ No newline at end of file
diff --git a/requirements.txt.bak b/requirements.txt.bak
new file mode 100644
index 0000000..8ddba70
--- /dev/null
+++ b/requirements.txt.bak
@@ -0,0 +1,16 @@
+librosa==0.9.1
+faster-whisper==0.9.0
+pydub==0.25.1
+wavmark==0.0.3
+numpy==1.22.0
+eng_to_ipa==0.0.2
+inflect==7.0.0
+unidecode==1.3.7
+whisper-timestamped==1.14.2
+openai
+python-dotenv
+pypinyin==0.50.0
+cn2an==0.5.22
+jieba==0.42.1
+gradio==3.48.0
+langid==1.1.6