openvoice_plugin

dreamvoice/.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1,2 @@
+from .api import DreamVoice
+from .plugin import DreamVoice_Plugin

dreamvoice/.ipynb_checkpoints/api-checkpoint.py

@@ -0,0 +1,295 @@
+import os
+import requests
+import yaml
+import torch
+import librosa
+import numpy as np
+import soundfile as sf
+from pathlib import Path
+from transformers import T5Tokenizer, T5EncoderModel
+from tqdm import tqdm
+from .src.vc_wrapper import ReDiffVC, DreamVC
+from .src.plugin_wrapper import DreamVG
+from .src.modules.speaker_encoder.encoder import inference as spk_encoder
+from .src.modules.BigVGAN.inference import load_model as load_vocoder
+from .src.feats.contentvec_hf import get_content_model, get_content
+
+
+class DreamVoice:
+    def __init__(self, config='dreamvc.yaml', mode='plugin', device='cuda', chunk_size=16):
+        # Initial setup
+        script_dir = Path(__file__).resolve().parent
+        config_path = script_dir / config
+
+        # Load configuration file
+        with open(config_path, 'r') as fp:
+            self.config = yaml.safe_load(fp)
+
+        self.script_dir = script_dir
+
+        # Ensure all checkpoints are downloaded
+        self._ensure_checkpoints_exist()
+
+        # Initialize attributes
+        self.device = device
+        self.sr = self.config['sample_rate']
+
+        # Load vocoder
+        vocoder_path = script_dir / self.config['vocoder_path']
+        self.hifigan, _ = load_vocoder(vocoder_path, device)
+        self.hifigan.eval()
+
+        # Load content model
+        self.content_model = get_content_model().to(device)
+
+        # Load tokenizer and text encoder
+        lm_path = self.config['lm_path']
+        self.tokenizer = T5Tokenizer.from_pretrained(lm_path)
+        self.text_encoder = T5EncoderModel.from_pretrained(lm_path).to(device).eval()
+
+        # Set mode
+        self.mode = mode
+        if mode == 'plugin':
+            self._init_plugin_mode()
+        elif mode == 'end2end':
+            self._init_end2end_mode()
+        else:
+            raise NotImplementedError("Select mode from 'plugin' and 'end2end'")
+
+        # chunk inputs to 10s clips
+        self.chunk_size = chunk_size * 50
+
+    def _ensure_checkpoints_exist(self):
+        checkpoints = [
+            ('vocoder_path', self.config.get('vocoder_url')),
+            ('vocoder_config_path', self.config.get('vocoder_config_url')),
+            ('speaker_path', self.config.get('speaker_url')),
+            ('dreamvc.ckpt_path', self.config.get('dreamvc', {}).get('ckpt_url')),
+            ('rediffvc.ckpt_path', self.config.get('rediffvc', {}).get('ckpt_url')),
+            ('dreamvg.ckpt_path', self.config.get('dreamvg', {}).get('ckpt_url'))
+        ]
+
+        for path_key, url in checkpoints:
+            local_path = self._get_local_path(path_key)
+            if not local_path.exists() and url:
+                print(f"Downloading {path_key} from {url}")
+                self._download_file(url, local_path)
+
+    def _get_local_path(self, path_key):
+        keys = path_key.split('.')
+        local_path = self.config
+        for key in keys:
+            local_path = local_path.get(key, {})
+        return self.script_dir / local_path
+
+    def _download_file(self, url, local_path):
+        try:
+            # Attempt to send a GET request to the URL
+            response = requests.get(url, stream=True)
+            response.raise_for_status()  # Ensure we raise an exception for HTTP errors
+        except requests.exceptions.RequestException as e:
+            # Log the error for debugging purposes
+            print(f"Error encountered: {e}")
+
+            # Development mode: prompt user for Hugging Face API key
+            user_input = input("Private checkpoint, please request authorization and enter your Hugging Face API key.")
+            self.hf_key = user_input if user_input else None
+
+            # Set headers if an API key is provided
+            headers = {'Authorization': f'Bearer {self.hf_key}'} if self.hf_key else {}
+
+            try:
+                # Attempt to send a GET request with headers in development mode
+                response = requests.get(url, stream=True, headers=headers)
+                response.raise_for_status()  # Ensure we raise an exception for HTTP errors
+            except requests.exceptions.RequestException as e:
+                # Log the error for debugging purposes
+                print(f"Error encountered in dev mode: {e}")
+                response = None  # Handle response accordingly in your code
+
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+
+        total_size = int(response.headers.get('content-length', 0))
+        block_size = 8192
+        t = tqdm(total=total_size, unit='iB', unit_scale=True)
+
+        with open(local_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=block_size):
+                t.update(len(chunk))
+                f.write(chunk)
+        t.close()
+
+    def _init_plugin_mode(self):
+        # Initialize ReDiffVC
+        self.dreamvc = ReDiffVC(
+            config_path=self.script_dir / self.config['rediffvc']['config_path'],
+            ckpt_path=self.script_dir / self.config['rediffvc']['ckpt_path'],
+            device=self.device
+        )
+
+        # Initialize DreamVG
+        self.dreamvg = DreamVG(
+            config_path=self.script_dir / self.config['dreamvg']['config_path'],
+            ckpt_path=self.script_dir / self.config['dreamvg']['ckpt_path'],
+            device=self.device
+        )
+
+        # Load speaker encoder
+        spk_encoder.load_model(self.script_dir / self.config['speaker_path'], self.device)
+        self.spk_encoder = spk_encoder
+        self.spk_embed_cache = None
+
+    def _init_end2end_mode(self):
+        # Initialize DreamVC
+        self.dreamvc = DreamVC(
+            config_path=self.script_dir / self.config['dreamvc']['config_path'],
+            ckpt_path=self.script_dir / self.config['dreamvc']['ckpt_path'],
+            device=self.device
+        )
+
+    def _load_content(self, audio_path):
+        content_audio, _ = librosa.load(audio_path, sr=16000)
+        # Calculate the required length to make it a multiple of 16*160
+        target_length = ((len(content_audio) + 16*160 - 1) // (16*160)) * (16*160)
+        # Pad with zeros if necessary
+        if len(content_audio) < target_length:
+            content_audio = np.pad(content_audio, (0, target_length - len(content_audio)), mode='constant')
+        content_audio = torch.tensor(content_audio).unsqueeze(0).to(self.device)
+        content_clip = get_content(self.content_model, content_audio)
+        return content_clip
+
+    def load_spk_embed(self, emb_path):
+        self.spk_embed_cache = torch.load(emb_path, map_location=self.device)
+
+    def save_spk_embed(self, emb_path):
+        assert self.spk_embed_cache is not None
+        torch.save(self.spk_embed_cache.cpu(), emb_path)
+
+    def save_audio(self, output_path, audio, sr):
+        sf.write(output_path, audio, samplerate=sr)
+
+    @torch.no_grad()
+    def genvc(self, content_audio, prompt,
+              prompt_guidance_scale=3, prompt_guidance_rescale=0.0,
+              prompt_ddim_steps=100, prompt_eta=1, prompt_random_seed=None,
+              vc_guidance_scale=3, vc_guidance_rescale=0.0,
+              vc_ddim_steps=50, vc_eta=1, vc_random_seed=None,
+             ):
+
+        content_clip = self._load_content(content_audio)
+
+        text_batch = self.tokenizer(prompt, max_length=32,
+                                    padding='max_length', truncation=True, return_tensors="pt")
+        text, text_mask = text_batch.input_ids.to(self.device), \
+            text_batch.attention_mask.to(self.device)
+        text = self.text_encoder(input_ids=text, attention_mask=text_mask)[0]
+
+        if self.mode == 'plugin':
+            spk_embed = self.dreamvg.inference([text, text_mask],
+                                               guidance_scale=prompt_guidance_scale,
+                                               guidance_rescale=prompt_guidance_rescale,
+                                               ddim_steps=prompt_ddim_steps, eta=prompt_eta,
+                                               random_seed=prompt_random_seed)
+            
+            B, L, D = content_clip.shape
+            gen_audio_chunks = []
+            num_chunks = (L + self.chunk_size - 1) // self.chunk_size
+            for i in range(num_chunks):
+                start_idx = i * self.chunk_size
+                end_idx = min((i + 1) * self.chunk_size, L)
+                content_clip_chunk = content_clip[:, start_idx:end_idx, :]
+                
+                gen_audio_chunk = self.dreamvc.inference(
+                        spk_embed, content_clip_chunk, None,
+                        guidance_scale=vc_guidance_scale,
+                        guidance_rescale=vc_guidance_rescale,
+                        ddim_steps=vc_ddim_steps,
+                        eta=vc_eta,
+                        random_seed=vc_random_seed)
+                    
+                gen_audio_chunks.append(gen_audio_chunk)
+            
+            gen_audio = torch.cat(gen_audio_chunks, dim=-1)
+
+            self.spk_embed_cache = spk_embed
+
+        elif self.mode == 'end2end':
+            B, L, D = content_clip.shape
+            gen_audio_chunks = []
+            num_chunks = (L + self.chunk_size - 1) // self.chunk_size
+            
+            for i in range(num_chunks):
+                start_idx = i * self.chunk_size
+                end_idx = min((i + 1) * self.chunk_size, L)
+                content_clip_chunk = content_clip[:, start_idx:end_idx, :]
+                
+                gen_audio_chunk = self.dreamvc.inference([text, text_mask], content_clip,
+                                                         guidance_scale=prompt_guidance_scale,
+                                                         guidance_rescale=prompt_guidance_rescale,
+                                                         ddim_steps=prompt_ddim_steps,
+                                                         eta=prompt_eta, random_seed=prompt_random_seed)
+                gen_audio_chunks.append(gen_audio_chunk)
+            
+            gen_audio = torch.cat(gen_audio_chunks, dim=-1)
+
+        else:
+            raise NotImplementedError("Select mode from 'plugin' and 'end2end'")
+
+        gen_audio = self.hifigan(gen_audio.squeeze(1))
+        gen_audio = gen_audio.cpu().numpy().squeeze(0).squeeze(0)
+
+        return gen_audio, self.sr
+
+    @torch.no_grad()
+    def simplevc(self, content_audio, speaker_audio=None, use_spk_cache=False,
+                 vc_guidance_scale=3, vc_guidance_rescale=0.0,
+                 vc_ddim_steps=50, vc_eta=1, vc_random_seed=None,
+                ):
+
+        assert self.mode == 'plugin'
+        if speaker_audio is not None:
+            speaker_audio, _ = librosa.load(speaker_audio, sr=16000)
+            speaker_audio = torch.tensor(speaker_audio).unsqueeze(0).to(self.device)
+            spk_embed = spk_encoder.embed_utterance_batch(speaker_audio)
+            self.spk_embed_cache = spk_embed
+        elif use_spk_cache:
+            assert self.spk_embed_cache is not None
+            spk_embed = self.spk_embed_cache
+        else:
+            raise NotImplementedError
+
+        content_clip = self._load_content(content_audio)
+
+        B, L, D = content_clip.shape
+        gen_audio_chunks = []
+        num_chunks = (L + self.chunk_size - 1) // self.chunk_size
+        for i in range(num_chunks):
+            start_idx = i * self.chunk_size
+            end_idx = min((i + 1) * self.chunk_size, L)
+            content_clip_chunk = content_clip[:, start_idx:end_idx, :]
+            
+            gen_audio_chunk = self.dreamvc.inference(
+                    spk_embed, content_clip_chunk, None,
+                    guidance_scale=vc_guidance_scale,
+                    guidance_rescale=vc_guidance_rescale,
+                    ddim_steps=vc_ddim_steps,
+                    eta=vc_eta,
+                    random_seed=vc_random_seed)
+                
+            gen_audio_chunks.append(gen_audio_chunk)
+        
+        gen_audio = torch.cat(gen_audio_chunks, dim=-1)
+
+        gen_audio = self.hifigan(gen_audio.squeeze(1))
+        gen_audio = gen_audio.cpu().numpy().squeeze(0).squeeze(0)
+
+        return gen_audio, self.sr
+
+
+if __name__ == '__main__':
+    dreamvoice = DreamVoice(config='dreamvc.yaml', mode='plugin', device='cuda')
+    content_audio = 'test.wav'
+    speaker_audio = 'speaker.wav'
+    prompt = 'young female voice, sounds young and cute'
+    gen_audio, sr = dreamvoice.genvc('test.wav', prompt)
+    dreamvoice.save_audio('debug.wav', gen_audio, sr)

dreamvoice/.ipynb_checkpoints/dreamvc-checkpoint.yaml

@@ -0,0 +1,27 @@
+version: 1.1
+
+sample_rate: 24000
+vocoder_path: 'ckpts/bigvgan_24k/g_01000000.pt'
+vocoder_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/ckpts/bigvgan_24k/g_01000000.pt'
+vocoder_config_path: 'ckpts/bigvgan_24k/config.json'
+vocoder_config_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/ckpts/bigvgan_24k/config.json'
+
+speaker_path: 'ckpts/spk_encoder/pretrained.pt'
+speaker_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/ckpts/spk_encoder/pretrained.pt'
+lm_path: 'google/flan-t5-base'
+
+dreamvc:
+  config_path: 'src/configs/diffvc_cross.yaml'
+  ckpt_path: 'ckpts/dreamvc_cross.pt'
+  ckpt_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/ckpts/dreamvc_cross.pt'
+
+rediffvc:
+  config_path: 'src/configs/diffvc_base.yaml'
+  ckpt_path: 'ckpts/dreamvc_base.pt'
+  ckpt_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/ckpts/dreamvc_base.pt'
+
+dreamvg:
+  config_path: 'src/configs/plugin_cross.yaml'
+  ckpt_path: 'ckpts/dreamvc_plugin.pt'
+  ckpt_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/ckpts/dreamvc_plugin.pt'
+

dreamvoice/.ipynb_checkpoints/openvoice_utils-checkpoint.py

@@ -0,0 +1,48 @@
+import os
+import torch
+import librosa
+from tqdm import tqdm
+from openvoice.mel_processing import spectrogram_torch
+from whisper_timestamped.transcribe import get_audio_tensor, get_vad_segments
+
+
+@torch.no_grad()
+def se_extractor(audio_path, vc):
+    # vad
+    SAMPLE_RATE = 16000
+    audio_vad = get_audio_tensor(audio_path)
+    segments = get_vad_segments(
+        audio_vad,
+        output_sample=True,
+        min_speech_duration=0.1,
+        min_silence_duration=1,
+        method="silero",
+    )
+    segments = [(seg["start"], seg["end"]) for seg in segments]
+    segments = [(float(s) / SAMPLE_RATE, float(e) / SAMPLE_RATE) for s,e in segments]
+
+    if len(segments) == 0:
+        segments = [(0, len(audio_vad)/SAMPLE_RATE)]
+        print(segments)
+
+    # spk
+    hps = vc.hps
+    device = vc.device
+    model = vc.model
+    gs = []
+
+    audio, sr = librosa.load(audio_path, sr=hps.data.sampling_rate)
+    audio = torch.tensor(audio).float().to(device)
+
+    for s, e in segments:
+        y = audio[int(hps.data.sampling_rate*s):int(hps.data.sampling_rate*e)]
+        y = y.to(device)
+        y = y.unsqueeze(0)
+        y = spectrogram_torch(y, hps.data.filter_length,
+                              hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,
+                              center=False).to(device)
+        g = model.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        gs.append(g.detach())
+
+    gs = torch.stack(gs).mean(0)
+    return gs.cpu()

dreamvoice/.ipynb_checkpoints/plugin-checkpoint.py

@@ -0,0 +1,128 @@
+import os
+import requests
+import yaml
+import torch
+import librosa
+import numpy as np
+import soundfile as sf
+from pathlib import Path
+from transformers import T5Tokenizer, T5EncoderModel
+from tqdm import tqdm
+from .src.plugin_wrapper import DreamVG
+
+
+class DreamVoice_Plugin:
+    def __init__(self, config='plugin.yaml', device='cuda'):
+        # Initial setup
+        script_dir = Path(__file__).resolve().parent
+        config_path = script_dir / config
+
+        # Load configuration file
+        with open(config_path, 'r') as fp:
+            self.config = yaml.safe_load(fp)
+
+        self.script_dir = script_dir
+
+        # Ensure all checkpoints are downloaded
+        self._ensure_checkpoints_exist()
+
+        # Initialize attributes
+        self.device = device
+
+        # Load tokenizer and text encoder
+        lm_path = self.config['lm_path']
+        self.tokenizer = T5Tokenizer.from_pretrained(lm_path)
+        self.text_encoder = T5EncoderModel.from_pretrained(lm_path).to(device).eval()
+        
+        self.dreamvg = DreamVG(
+            config_path=self.script_dir / self.config['dreamvg']['config_path'],
+            ckpt_path=self.script_dir / self.config['dreamvg']['ckpt_path'],
+            device=self.device
+        
+        )
+    def _ensure_checkpoints_exist(self):
+        checkpoints = [
+            ('dreamvg.ckpt_path', self.config.get('dreamvg', {}).get('ckpt_url'))
+        ]
+
+        for path_key, url in checkpoints:
+            local_path = self._get_local_path(path_key)
+            if not local_path.exists() and url:
+                print(f"Downloading {path_key} from {url}")
+                self._download_file(url, local_path)
+
+    def _get_local_path(self, path_key):
+        keys = path_key.split('.')
+        local_path = self.config
+        for key in keys:
+            local_path = local_path.get(key, {})
+        return self.script_dir / local_path
+
+    def _download_file(self, url, local_path):
+        try:
+            # Attempt to send a GET request to the URL
+            response = requests.get(url, stream=True)
+            response.raise_for_status()  # Ensure we raise an exception for HTTP errors
+        except requests.exceptions.RequestException as e:
+            # Log the error for debugging purposes
+            print(f"Error encountered: {e}")
+
+            # Development mode: prompt user for Hugging Face API key
+            user_input = input("Private checkpoint, please request authorization and enter your Hugging Face API key.")
+            self.hf_key = user_input if user_input else None
+
+            # Set headers if an API key is provided
+            headers = {'Authorization': f'Bearer {self.hf_key}'} if self.hf_key else {}
+
+            try:
+                # Attempt to send a GET request with headers in development mode
+                response = requests.get(url, stream=True, headers=headers)
+                response.raise_for_status()  # Ensure we raise an exception for HTTP errors
+            except requests.exceptions.RequestException as e:
+                # Log the error for debugging purposes
+                print(f"Error encountered in dev mode: {e}")
+                response = None  # Handle response accordingly in your code
+
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+
+        total_size = int(response.headers.get('content-length', 0))
+        block_size = 8192
+        t = tqdm(total=total_size, unit='iB', unit_scale=True)
+
+        with open(local_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=block_size):
+                t.update(len(chunk))
+                f.write(chunk)
+        t.close()
+
+    def _init_plugin_mode(self):
+        # Initialize DreamVG
+        self.dreamvg = DreamVG(
+            config_path=self.script_dir / self.config['dreamvg']['config_path'],
+            ckpt_path=self.script_dir / self.config['dreamvg']['ckpt_path'],
+            device=self.device
+        )
+
+        # Load speaker encoder
+        spk_encoder.load_model(self.script_dir / self.config['speaker_path'], self.device)
+        self.spk_encoder = spk_encoder
+        self.spk_embed_cache = None
+
+
+    @torch.no_grad()
+    def gen_spk(self, prompt,
+                prompt_guidance_scale=3, prompt_guidance_rescale=0.0,
+                prompt_ddim_steps=100, prompt_eta=1, prompt_random_seed=None,):
+        
+        text_batch = self.tokenizer(prompt, max_length=32,
+                                    padding='max_length', truncation=True, return_tensors="pt")
+        text, text_mask = text_batch.input_ids.to(self.device), \
+            text_batch.attention_mask.to(self.device)
+        text = self.text_encoder(input_ids=text, attention_mask=text_mask)[0]
+        
+        spk_embed = self.dreamvg.inference([text, text_mask],
+                                           guidance_scale=prompt_guidance_scale,
+                                           guidance_rescale=prompt_guidance_rescale,
+                                           ddim_steps=prompt_ddim_steps, eta=prompt_eta,
+                                           random_seed=prompt_random_seed)
+        return spk_embed

dreamvoice/.ipynb_checkpoints/plugin-checkpoint.yaml

@@ -0,0 +1,8 @@
+version: 1.1
+
+lm_path: 'google/flan-t5-base'
+
+dreamvg:
+  config_path: 'src/configs/plugin_cross_openvoice.yaml'
+  ckpt_path: 'plugin_ckpts/openvoice_v2.pt'
+  ckpt_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/plugin_ckpts/openvoice_v2.pt'

dreamvoice/__init__.py

@@ -1 +1,2 @@
-from .api import DreamVoice
+from .api import DreamVoice
+from .plugin import DreamVoice_Plugin

dreamvoice/dreamvc.yaml

@@ -1,4 +1,4 @@
-version: 1.0
+version: 1.1
 
 sample_rate: 24000
 vocoder_path: 'ckpts/bigvgan_24k/g_01000000.pt'
@@ -24,3 +24,4 @@ dreamvg:
   config_path: 'src/configs/plugin_cross.yaml'
   ckpt_path: 'ckpts/dreamvc_plugin.pt'
   ckpt_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/ckpts/dreamvc_plugin.pt'
+

dreamvoice/openvoice_utils.py

@@ -0,0 +1,48 @@
+import os
+import torch
+import librosa
+from tqdm import tqdm
+from openvoice.mel_processing import spectrogram_torch
+from whisper_timestamped.transcribe import get_audio_tensor, get_vad_segments
+
+
+@torch.no_grad()
+def se_extractor(audio_path, vc):
+    # vad
+    SAMPLE_RATE = 16000
+    audio_vad = get_audio_tensor(audio_path)
+    segments = get_vad_segments(
+        audio_vad,
+        output_sample=True,
+        min_speech_duration=0.1,
+        min_silence_duration=1,
+        method="silero",
+    )
+    segments = [(seg["start"], seg["end"]) for seg in segments]
+    segments = [(float(s) / SAMPLE_RATE, float(e) / SAMPLE_RATE) for s,e in segments]
+
+    if len(segments) == 0:
+        segments = [(0, len(audio_vad)/SAMPLE_RATE)]
+        print(segments)
+
+    # spk
+    hps = vc.hps
+    device = vc.device
+    model = vc.model
+    gs = []
+
+    audio, sr = librosa.load(audio_path, sr=hps.data.sampling_rate)
+    audio = torch.tensor(audio).float().to(device)
+
+    for s, e in segments:
+        y = audio[int(hps.data.sampling_rate*s):int(hps.data.sampling_rate*e)]
+        y = y.to(device)
+        y = y.unsqueeze(0)
+        y = spectrogram_torch(y, hps.data.filter_length,
+                              hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,
+                              center=False).to(device)
+        g = model.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        gs.append(g.detach())
+
+    gs = torch.stack(gs).mean(0)
+    return gs.cpu()

dreamvoice/plugin.py

@@ -0,0 +1,128 @@
+import os
+import requests
+import yaml
+import torch
+import librosa
+import numpy as np
+import soundfile as sf
+from pathlib import Path
+from transformers import T5Tokenizer, T5EncoderModel
+from tqdm import tqdm
+from .src.plugin_wrapper import DreamVG
+
+
+class DreamVoice_Plugin:
+    def __init__(self, config='plugin.yaml', device='cuda'):
+        # Initial setup
+        script_dir = Path(__file__).resolve().parent
+        config_path = script_dir / config
+
+        # Load configuration file
+        with open(config_path, 'r') as fp:
+            self.config = yaml.safe_load(fp)
+
+        self.script_dir = script_dir
+
+        # Ensure all checkpoints are downloaded
+        self._ensure_checkpoints_exist()
+
+        # Initialize attributes
+        self.device = device
+
+        # Load tokenizer and text encoder
+        lm_path = self.config['lm_path']
+        self.tokenizer = T5Tokenizer.from_pretrained(lm_path)
+        self.text_encoder = T5EncoderModel.from_pretrained(lm_path).to(device).eval()
+        
+        self.dreamvg = DreamVG(
+            config_path=self.script_dir / self.config['dreamvg']['config_path'],
+            ckpt_path=self.script_dir / self.config['dreamvg']['ckpt_path'],
+            device=self.device
+        
+        )
+    def _ensure_checkpoints_exist(self):
+        checkpoints = [
+            ('dreamvg.ckpt_path', self.config.get('dreamvg', {}).get('ckpt_url'))
+        ]
+
+        for path_key, url in checkpoints:
+            local_path = self._get_local_path(path_key)
+            if not local_path.exists() and url:
+                print(f"Downloading {path_key} from {url}")
+                self._download_file(url, local_path)
+
+    def _get_local_path(self, path_key):
+        keys = path_key.split('.')
+        local_path = self.config
+        for key in keys:
+            local_path = local_path.get(key, {})
+        return self.script_dir / local_path
+
+    def _download_file(self, url, local_path):
+        try:
+            # Attempt to send a GET request to the URL
+            response = requests.get(url, stream=True)
+            response.raise_for_status()  # Ensure we raise an exception for HTTP errors
+        except requests.exceptions.RequestException as e:
+            # Log the error for debugging purposes
+            print(f"Error encountered: {e}")
+
+            # Development mode: prompt user for Hugging Face API key
+            user_input = input("Private checkpoint, please request authorization and enter your Hugging Face API key.")
+            self.hf_key = user_input if user_input else None
+
+            # Set headers if an API key is provided
+            headers = {'Authorization': f'Bearer {self.hf_key}'} if self.hf_key else {}
+
+            try:
+                # Attempt to send a GET request with headers in development mode
+                response = requests.get(url, stream=True, headers=headers)
+                response.raise_for_status()  # Ensure we raise an exception for HTTP errors
+            except requests.exceptions.RequestException as e:
+                # Log the error for debugging purposes
+                print(f"Error encountered in dev mode: {e}")
+                response = None  # Handle response accordingly in your code
+
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+
+        total_size = int(response.headers.get('content-length', 0))
+        block_size = 8192
+        t = tqdm(total=total_size, unit='iB', unit_scale=True)
+
+        with open(local_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=block_size):
+                t.update(len(chunk))
+                f.write(chunk)
+        t.close()
+
+    def _init_plugin_mode(self):
+        # Initialize DreamVG
+        self.dreamvg = DreamVG(
+            config_path=self.script_dir / self.config['dreamvg']['config_path'],
+            ckpt_path=self.script_dir / self.config['dreamvg']['ckpt_path'],
+            device=self.device
+        )
+
+        # Load speaker encoder
+        spk_encoder.load_model(self.script_dir / self.config['speaker_path'], self.device)
+        self.spk_encoder = spk_encoder
+        self.spk_embed_cache = None
+
+
+    @torch.no_grad()
+    def gen_spk(self, prompt,
+                prompt_guidance_scale=3, prompt_guidance_rescale=0.0,
+                prompt_ddim_steps=100, prompt_eta=1, prompt_random_seed=None,):
+        
+        text_batch = self.tokenizer(prompt, max_length=32,
+                                    padding='max_length', truncation=True, return_tensors="pt")
+        text, text_mask = text_batch.input_ids.to(self.device), \
+            text_batch.attention_mask.to(self.device)
+        text = self.text_encoder(input_ids=text, attention_mask=text_mask)[0]
+        
+        spk_embed = self.dreamvg.inference([text, text_mask],
+                                           guidance_scale=prompt_guidance_scale,
+                                           guidance_rescale=prompt_guidance_rescale,
+                                           ddim_steps=prompt_ddim_steps, eta=prompt_eta,
+                                           random_seed=prompt_random_seed)
+        return spk_embed

dreamvoice/plugin.yaml

@@ -0,0 +1,8 @@
+version: 1.1
+
+lm_path: 'google/flan-t5-base'
+
+dreamvg:
+  config_path: 'src/configs/plugin_cross_openvoice.yaml'
+  ckpt_path: 'plugin_ckpts/openvoice_v2.pt'
+  ckpt_url: 'https://huggingface.co/myshell-ai/DreamVoice/resolve/main/dreamvoice/plugin_ckpts/openvoice_v2.pt'

dreamvoice/plugin_ckpts/openvoice_v2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08eae6c5b5da6438451589f3657b3bfeedc3b80bed52948f9845b7d70e989ab1
+size 104892189

dreamvoice/src/.ipynb_checkpoints/plugin_wrapper-checkpoint.py

@@ -0,0 +1,76 @@
+import yaml
+import torch
+from diffusers import DDIMScheduler
+from .model.p2e_cross import P2E_Cross
+from .utils import scale_shift, scale_shift_re, rescale_noise_cfg
+
+
+class DreamVG(object):
+    def __init__(self,
+                 config_path='configs/plugin_cross.yaml',
+                 ckpt_path='../ckpts/dreamvc_plugin.pt',
+                 device='cpu'):
+
+        with open(config_path, 'r') as fp:
+            config = yaml.safe_load(fp)
+
+        self.device = device
+        self.model = P2E_Cross(config['model']).to(device)
+        self.model.load_state_dict(torch.load(ckpt_path)['model'])
+        self.model.eval()
+
+        noise_scheduler = DDIMScheduler(num_train_timesteps=config['scheduler']['num_train_steps'],
+                                        beta_start=config['scheduler']['beta_start'],
+                                        beta_end=config['scheduler']['beta_end'],
+                                        rescale_betas_zero_snr=True,
+                                        timestep_spacing="trailing",
+                                        clip_sample=False,
+                                        prediction_type='v_prediction')
+        self.noise_scheduler = noise_scheduler
+        self.scale = config['scheduler']['scale']
+        self.shift = config['scheduler']['shift']
+        self.spk_shape = config['model']['unet']['in_channels']
+
+    @torch.no_grad()
+    def inference(self, text,
+                  guidance_scale=5, guidance_rescale=0.7,
+                  ddim_steps=50, eta=1, random_seed=2023,
+                 ):
+        text, text_mask = text
+        self.model.eval()
+
+        gen_shape = (1, self.spk_shape)
+        
+        if random_seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(random_seed)
+        else:
+            generator = torch.Generator(device=self.device)
+            generator.seed()
+
+        self.noise_scheduler.set_timesteps(ddim_steps)
+    
+        # init noise
+        noise = torch.randn(gen_shape, generator=generator, device=self.device)
+        latents = noise
+    
+        for t in self.noise_scheduler.timesteps:
+            latents = self.noise_scheduler.scale_model_input(latents, t)
+    
+            if guidance_scale:
+                output_text = self.model(latents, t, text, text_mask, train_cfg=False)
+                output_uncond = self.model(latents, t, text, text_mask, train_cfg=True, cfg_prob=1.0)
+    
+                output_pred = output_uncond + guidance_scale * (output_text - output_uncond)
+                if guidance_rescale > 0.0:
+                    output_pred = rescale_noise_cfg(output_pred, output_text,
+                                                    guidance_rescale=guidance_rescale)
+            else:
+                output_pred = self.model(latents, t, text, text_mask, train_cfg=False)
+    
+            latents = self.noise_scheduler.step(model_output=output_pred, timestep=t, sample=latents,
+                                                eta=eta, generator=generator).prev_sample
+    
+        # pred = reverse_minmax_norm_diff(latents, vmin=0.0, vmax=0.5)
+        pred = scale_shift_re(latents, 1/self.scale, self.shift)
+        # pred = torch.clip(pred, min=0.0, max=0.5)
+        return pred

dreamvoice/src/.ipynb_checkpoints/vc_wrapper-checkpoint.py

@@ -0,0 +1,144 @@
+import yaml
+import torch
+from diffusers import DDIMScheduler
+from .model.model import DiffVC
+from .model.model_cross import DiffVC_Cross
+from .utils import scale_shift, scale_shift_re, rescale_noise_cfg
+
+
+class ReDiffVC(object):
+    def __init__(self,
+                 config_path='configs/diffvc_base.yaml',
+                 ckpt_path='../ckpts/dreamvc_base.pt',
+                 device='cpu'):
+
+        with open(config_path, 'r') as fp:
+            config = yaml.safe_load(fp)
+
+        self.device = device
+        self.model = DiffVC(config['model']).to(device)
+        self.model.load_state_dict(torch.load(ckpt_path)['model'])
+        self.model.eval()
+
+        noise_scheduler = DDIMScheduler(num_train_timesteps=config['scheduler']['num_train_steps'],
+                                        beta_start=config['scheduler']['beta_start'],
+                                        beta_end=config['scheduler']['beta_end'],
+                                        rescale_betas_zero_snr=True,
+                                        timestep_spacing="trailing",
+                                        clip_sample=False,
+                                        prediction_type='v_prediction')
+        self.noise_scheduler = noise_scheduler
+        self.scale = config['scheduler']['scale']
+        self.shift = config['scheduler']['shift']
+        self.melshape = config['model']['unet']['sample_size'][0]
+
+    @torch.no_grad()
+    def inference(self,
+                  spk_embed, content_clip, f0_clip=None,
+                  guidance_scale=3, guidance_rescale=0.7,
+                  ddim_steps=50, eta=1, random_seed=2023):
+
+        self.model.eval()
+        if random_seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(random_seed)
+        else:
+            generator = torch.Generator(device=self.device)
+            generator.seed()
+
+        self.noise_scheduler.set_timesteps(ddim_steps)
+
+        # init noise
+        gen_shape = (1, 1, self.melshape, content_clip.shape[-2])
+        noise = torch.randn(gen_shape, generator=generator, device=self.device)
+        latents = noise
+
+        for t in self.noise_scheduler.timesteps:
+            latents = self.noise_scheduler.scale_model_input(latents, t)
+
+            if guidance_scale:
+                output_text = self.model(latents, t, content_clip, spk_embed, f0_clip, train_cfg=False)
+                output_uncond = self.model(latents, t, content_clip, spk_embed, f0_clip, train_cfg=True,
+                                           speaker_cfg=1.0, pitch_cfg=0.0)
+
+                output_pred = output_uncond + guidance_scale * (output_text - output_uncond)
+                if guidance_rescale > 0.0:
+                    output_pred = rescale_noise_cfg(output_pred, output_text,
+                                                    guidance_rescale=guidance_rescale)
+            else:
+                output_pred = self.model(latents, t, content_clip, spk_embed, f0_clip, train_cfg=False)
+
+            latents = self.noise_scheduler.step(model_output=output_pred, timestep=t, sample=latents,
+                                                eta=eta, generator=generator).prev_sample
+
+        pred = scale_shift_re(latents, scale=1/self.scale, shift=self.shift)
+        return pred
+
+
+class DreamVC(object):
+    def __init__(self,
+                 config_path='configs/diffvc_cross.yaml',
+                 ckpt_path='../ckpts/dreamvc_cross.pt',
+                 device='cpu'):
+
+        with open(config_path, 'r') as fp:
+            config = yaml.safe_load(fp)
+
+        self.device = device
+        self.model = DiffVC_Cross(config['model']).to(device)
+        self.model.load_state_dict(torch.load(ckpt_path)['model'])
+        self.model.eval()
+
+        noise_scheduler = DDIMScheduler(num_train_timesteps=config['scheduler']['num_train_steps'],
+                                        beta_start=config['scheduler']['beta_start'],
+                                        beta_end=config['scheduler']['beta_end'],
+                                        rescale_betas_zero_snr=True,
+                                        timestep_spacing="trailing",
+                                        clip_sample=False,
+                                        prediction_type='v_prediction')
+        self.noise_scheduler = noise_scheduler
+        self.scale = config['scheduler']['scale']
+        self.shift = config['scheduler']['shift']
+        self.melshape = config['model']['unet']['sample_size'][0]
+
+    @torch.no_grad()
+    def inference(self,
+                  text, content_clip, f0_clip=None,
+                  guidance_scale=3, guidance_rescale=0.7,
+                  ddim_steps=50, eta=1, random_seed=2023):
+
+        text, text_mask = text
+        self.model.eval()
+        if random_seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(random_seed)
+        else:
+            generator = torch.Generator(device=self.device)
+            generator.seed()
+
+        self.noise_scheduler.set_timesteps(ddim_steps)
+
+        # init noise
+        gen_shape = (1, 1, self.melshape, content_clip.shape[-2])
+        noise = torch.randn(gen_shape, generator=generator, device=self.device)
+        latents = noise
+
+        for t in self.noise_scheduler.timesteps:
+            latents = self.noise_scheduler.scale_model_input(latents, t)
+
+            if guidance_scale:
+                output_text = self.model(latents, t, content_clip, text, text_mask, f0_clip, train_cfg=False)
+                output_uncond = self.model(latents, t, content_clip, text, text_mask, f0_clip, train_cfg=True,
+                                           speaker_cfg=1.0, pitch_cfg=0.0)
+
+                output_pred = output_uncond + guidance_scale * (output_text - output_uncond)
+                if guidance_rescale > 0.0:
+                    output_pred = rescale_noise_cfg(output_pred, output_text,
+                                                    guidance_rescale=guidance_rescale)
+            else:
+                output_pred = self.model(latents, t, content_clip, text, text_mask, f0_clip, train_cfg=False)
+
+            latents = self.noise_scheduler.step(model_output=output_pred, timestep=t, sample=latents,
+                                                eta=eta, generator=generator).prev_sample
+
+        pred = scale_shift_re(latents, scale=1/self.scale, shift=self.shift)
+        return pred
+

dreamvoice/src/configs/.ipynb_checkpoints/plugin_cross-checkpoint.yaml

@@ -0,0 +1,39 @@
+version: 1.0
+
+system: "cross"
+
+model:
+  cls_embedding:
+    content_dim: 768
+    content_hidden: 256
+
+  unet:  
+    sample_size: [1, 1]
+    in_channels: 256
+    out_channels: 256
+    layers_per_block: 2
+    block_out_channels: [256]
+    down_block_types:
+      [
+        "CrossAttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "CrossAttnUpBlock2D",
+      ]
+    attention_head_dim: 32
+    cross_attention_dim: 768
+
+scheduler:
+  num_train_steps: 1000
+  beta_schedule: 'linear'
+  beta_start: 0.0001
+  beta_end: 0.02
+  num_infer_steps: 50
+  rescale_betas_zero_snr: true
+  timestep_spacing: "trailing"
+  clip_sample: false
+  prediction_type: 'v_prediction'
+  scale: 0.05 
+  shift: -0.035
+    

dreamvoice/src/configs/.ipynb_checkpoints/plugin_cross_openvoice-checkpoint.yaml

@@ -0,0 +1,39 @@
+version: 1.0
+
+system: "cross"
+
+model:
+  cls_embedding:
+    content_dim: 768
+    content_hidden: 256
+
+  unet:  
+    sample_size: [1, 1]
+    in_channels: 256
+    out_channels: 256
+    layers_per_block: 2
+    block_out_channels: [256]
+    down_block_types:
+      [
+        "CrossAttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "CrossAttnUpBlock2D",
+      ]
+    attention_head_dim: 32
+    cross_attention_dim: 768
+
+scheduler:
+  num_train_steps: 1000
+  beta_schedule: 'linear'
+  beta_start: 0.0001
+  beta_end: 0.02
+  num_infer_steps: 50
+  rescale_betas_zero_snr: true
+  timestep_spacing: "trailing"
+  clip_sample: false
+  prediction_type: 'v_prediction'
+  scale: 1.0
+  shift: 0.0
+    

dreamvoice/src/configs/plugin_cross_openvoice.yaml

@@ -0,0 +1,39 @@
+version: 1.0
+
+system: "cross"
+
+model:
+  cls_embedding:
+    content_dim: 768
+    content_hidden: 256
+
+  unet:  
+    sample_size: [1, 1]
+    in_channels: 256
+    out_channels: 256
+    layers_per_block: 2
+    block_out_channels: [256]
+    down_block_types:
+      [
+        "CrossAttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "CrossAttnUpBlock2D",
+      ]
+    attention_head_dim: 32
+    cross_attention_dim: 768
+
+scheduler:
+  num_train_steps: 1000
+  beta_schedule: 'linear'
+  beta_start: 0.0001
+  beta_end: 0.02
+  num_infer_steps: 50
+  rescale_betas_zero_snr: true
+  timestep_spacing: "trailing"
+  clip_sample: false
+  prediction_type: 'v_prediction'
+  scale: 1.0
+  shift: 0.0
+