Initial commit

dreamvoice/.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1 @@
+from .api import DreamVoice

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.7,
+              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.7,
+                 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,26 @@
+version: 1.0
+
+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/__init__.py

@@ -0,0 +1 @@
+from .api import DreamVoice

dreamvoice/api.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.7,
+              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.7,
+                 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/ckpts/bigvgan_24k/config.json

@@ -0,0 +1,44 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 32,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.999,
+    "seed": 1234,
+
+    "upsample_rates": [10,6,4,2],
+    "upsample_kernel_sizes": [20,12,8,4],
+    "upsample_initial_channel": 512,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "segment_size": 12000,
+    "num_mels": 128,
+    "n_fft": 1920,
+    "hop_size": 480,
+    "win_size": 1920,
+
+    "sampling_rate": 24000,
+
+    "fmin": 0,
+    "fmax": 12000,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

dreamvoice/ckpts/bigvgan_24k/g_01000000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:683a7baafedda8ec2fd2409deff61bd58ae66fbf10630550a17fcfed6f728977
+size 58405452

dreamvoice/ckpts/dreamvc_base.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5abe034bf590e2ce0405c66e950dc61f041629731e959cb09e2009688cd1254c
+size 300117179

dreamvoice/ckpts/dreamvc_cross.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87b4eb1e62b1bf4e157edc2766b9b4461c0be0f7d98a970d6b087f3797c35920
+size 451974443

dreamvoice/ckpts/dreamvc_plugin.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2396f6b96e9057e73e20eee173d7aaded6b5eb70745a9f5282999c0ea9a4d848
+size 104892440

dreamvoice/ckpts/spk_encoder/pretrained.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39373b86598fa3da9fcddee6142382efe09777e8d37dc9c0561f41f0070f134e
+size 17090379

dreamvoice/dreamvc.yaml

@@ -0,0 +1,26 @@
+version: 1.0
+
+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/src/.ipynb_checkpoints/extract_features-checkpoint.py

@@ -0,0 +1,103 @@
+import os
+import torch
+import librosa
+import numpy as np
+import soundfile as sf
+import pandas as pd
+# from feats.hubert_model import get_soft_model, get_hubert_soft_content
+from feats.contentvec_hf import get_content_model, get_content
+# from modules.speaker_encoder.encoder import inference as spk_encoder
+# from pathlib import Path
+from tqdm import tqdm
+from multiprocessing import Process
+import pyworld as pw
+
+
+def resample_save(infolder, audio_path, model,
+                  audio_sr=24000, content_sr=16000, min_length=1.92,
+                  content_resolution=50,
+                  save_path='features'):
+    if os.path.isfile(save_path + '/' + 'audio_24k/' + audio_path) is False:
+        audio, sr = librosa.load(infolder + audio_path, sr=content_sr)
+        final_length = audio.shape[-1] // (content_sr / content_resolution) * (content_sr / content_resolution)
+        # final_length = final_length / content_sr
+
+        length = max(round(min_length*content_sr), round(final_length))
+        assert length % 10 == 0
+        audio = audio[:length]
+        audio_save = np.zeros(length, dtype=audio.dtype)
+        audio_save[:audio.shape[-1]] = audio[:audio.shape[-1]]
+
+        # content = get_hubert_soft_content(model, torch.tensor(audio_save).unsqueeze(0))
+        content = get_content(model, torch.tensor(audio_save).unsqueeze(0))
+        content = content.cpu()
+        os.makedirs(os.path.dirname(save_path + '/' + 'content/' + audio_path), exist_ok=True)
+        torch.save(content, save_path + '/' + 'content/' + audio_path+'.pt')
+        # print(audio_save.shape)
+        # print(content.shape)
+        os.makedirs(os.path.dirname(save_path + '/' + 'audio_16k/' + audio_path), exist_ok=True)
+        sf.write(save_path + '/' + 'audio_16k/' + audio_path, audio_save, int(sr))
+        # print(save_path + '/' + 'audio_16k/' + audio_path)
+
+        audio, sr = librosa.load(infolder + audio_path, sr=audio_sr)
+        length = max(round(min_length*audio_sr), round(final_length/content_sr*audio_sr))
+        assert length % 10 == 0
+        audio = audio[:length]
+        audio_save = np.zeros(length, dtype=audio.dtype)
+        audio_save[:audio.shape[-1]] = audio[:audio.shape[-1]]
+        # print(audio_save.shape)
+        os.makedirs(os.path.dirname(save_path + '/' + 'audio_24k/' + audio_path), exist_ok=True)
+        sf.write(save_path + '/' + 'audio_24k/' + audio_path, audio_save, int(sr))
+
+
+def extract_f0(in_folder, audio_path, save_path):
+    audio, sr = librosa.load(in_folder + audio_path, sr=None)
+    assert sr == 16000
+    if os.path.isfile(save_path + '/' + 'f0/' + audio_path + '.pt') is False:
+        # wav = audio
+        # wav = np.pad(wav, int((1024-320)/2), mode='reflect')
+        # f0_, _, _ = librosa.pyin(wav, frame_length=1024, hop_length=320, center=False, sr=sr,
+        #                         fmin=librosa.note_to_hz('C2'),
+        #                         fmax=librosa.note_to_hz('C6'))
+
+        _f0, t = pw.dio(audio.astype(np.float64), sr, frame_period=320 / sr * 1000)
+        f0 = pw.stonemask(audio.astype(np.float64), _f0, t, sr)[:-1]
+
+        f0 = np.nan_to_num(f0)
+        os.makedirs(os.path.dirname(save_path + '/' + 'f0/' + audio_path), exist_ok=True)
+        # print(save_path + '/' + 'f0/' + audio_path + '.pt')
+        torch.save(torch.tensor(f0), save_path + '/' + 'f0/' + audio_path + '.pt')
+
+
+def chunks(arr, m):
+    result = [[] for i in range(m)]
+    for i in range(len(arr)):
+        result[i%m].append(arr[i])
+    return result
+
+
+def extract_f0_main(in_folder, audio_paths, save_path):
+    for audio_path in tqdm(audio_paths):
+        extract_f0(in_folder, audio_path, save_path)
+
+
+if __name__ == '__main__':
+    df = pd.read_csv('../test_data/vc_meta.csv')
+    # model = get_soft_model('../pre_ckpts/hubert_soft.pt').to('cuda')
+    model = get_content_model().to('cuda')
+    # # spk_encoder.load_model(Path('ckpts/spk_encoder/pretrained.pt'), device="cuda")
+    for i in tqdm(range(len(df))):
+        row = df.iloc[i]
+        in_path = row['path']
+        resample_save('../test_data/', in_path, model, save_path='../features/')
+
+    in_folder = '../features/audio_16k/'
+    audio_files = list(df['path'])
+    save_path = '../features/'
+    cores = 6
+
+    subsets = chunks(audio_files, cores)
+
+    for subset in subsets:
+        t = Process(target=extract_f0_main, args=(in_folder, subset, save_path))
+        t.start()

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/diffvc_base-checkpoint.yaml

@@ -0,0 +1,47 @@
+version: 1.0
+
+system: "base"
+
+model:
+  cls_embedding:
+    speaker_dim: 256
+    feature_dim: 512
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: false
+
+  unet:  
+    sample_size: [128, 256]
+    in_channels: 257
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [128, 256, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "DownBlock2D",
+        "AttnDownBlock2D",
+        "AttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "AttnUpBlock2D",
+        "AttnUpBlock2D",
+        "UpBlock2D",
+        "UpBlock2D"
+      ]
+    attention_head_dim: 32
+    class_embed_type: 'identity'
+
+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: 2.75
+  shift: 5.80

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

@@ -0,0 +1,34 @@
+version: 1.0
+
+system: "base"
+
+diffwrap:
+  cls_embedding:
+    speaker_dim: 256
+    feature_dim: 512
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: true
+    pitch_dim: 1
+    pitch_hidden: 128
+    
+  unet:  
+    sample_size: [128, 256]
+    in_channels: 385
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [256, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "AttnDownBlock2D",
+        "AttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "AttnUpBlock2D",
+        "AttnUpBlock2D",
+        "UpBlock2D"
+      ]
+    attention_head_dim: 32
+    class_embed_type: 'identity'

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

@@ -0,0 +1,45 @@
+version: 1.0
+
+system: "cross"
+
+model:
+  cls_embedding:
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: false
+
+  unet:  
+    sample_size: [128, 256]
+    in_channels: 257
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [128, 256, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "DownBlock2D",
+        "CrossAttnDownBlock2D",
+        "CrossAttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "CrossAttnUpBlock2D",
+        "CrossAttnUpBlock2D",
+        "UpBlock2D",
+        "UpBlock2D",
+      ]
+    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: 2.75
+  shift: 5.80

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

@@ -0,0 +1,33 @@
+version: 1.0
+
+system: "cross"
+
+diffwrap:
+  cls_embedding:
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: true
+    pitch_dim: 1
+    pitch_hidden: 128
+
+  unet:  
+    sample_size: [100, 256]
+    in_channels: 385
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [128, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "CrossAttnDownBlock2D",
+        "CrossAttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "CrossAttnUpBlock2D",
+        "CrossAttnUpBlock2D",
+        "UpBlock2D",
+      ]
+    attention_head_dim: 32
+    cross_attention_dim: 768
+    

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/diffvc_base.yaml

@@ -0,0 +1,47 @@
+version: 1.0
+
+system: "base"
+
+model:
+  cls_embedding:
+    speaker_dim: 256
+    feature_dim: 512
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: false
+
+  unet:  
+    sample_size: [128, 256]
+    in_channels: 257
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [128, 256, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "DownBlock2D",
+        "AttnDownBlock2D",
+        "AttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "AttnUpBlock2D",
+        "AttnUpBlock2D",
+        "UpBlock2D",
+        "UpBlock2D"
+      ]
+    attention_head_dim: 32
+    class_embed_type: 'identity'
+
+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: 2.75
+  shift: 5.80

dreamvoice/src/configs/diffvc_base_pitch.yaml

@@ -0,0 +1,34 @@
+version: 1.0
+
+system: "base"
+
+diffwrap:
+  cls_embedding:
+    speaker_dim: 256
+    feature_dim: 512
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: true
+    pitch_dim: 1
+    pitch_hidden: 128
+    
+  unet:  
+    sample_size: [128, 256]
+    in_channels: 385
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [128, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "AttnDownBlock2D",
+        "AttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "AttnUpBlock2D",
+        "AttnUpBlock2D",
+        "UpBlock2D"
+      ]
+    attention_head_dim: 32
+    class_embed_type: 'identity'

dreamvoice/src/configs/diffvc_cross.yaml

@@ -0,0 +1,45 @@
+version: 1.0
+
+system: "cross"
+
+model:
+  cls_embedding:
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: false
+
+  unet:  
+    sample_size: [128, 256]
+    in_channels: 257
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [128, 256, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "DownBlock2D",
+        "CrossAttnDownBlock2D",
+        "CrossAttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "CrossAttnUpBlock2D",
+        "CrossAttnUpBlock2D",
+        "UpBlock2D",
+        "UpBlock2D",
+      ]
+    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: 2.75
+  shift: 5.80

dreamvoice/src/configs/diffvc_cross_pitch.yaml

@@ -0,0 +1,33 @@
+version: 1.0
+
+system: "cross"
+
+diffwrap:
+  cls_embedding:
+    content_dim: 768
+    content_hidden: 256
+    use_pitch: true
+    pitch_dim: 1
+    pitch_hidden: 128
+
+  unet:  
+    sample_size: [100, 256]
+    in_channels: 385
+    out_channels: 1
+    layers_per_block: 2
+    block_out_channels: [128, 256, 512]
+    down_block_types:
+      [
+        "DownBlock2D",
+        "CrossAttnDownBlock2D",
+        "CrossAttnDownBlock2D",
+      ]
+    up_block_types:
+      [
+        "CrossAttnUpBlock2D",
+        "CrossAttnUpBlock2D",
+        "UpBlock2D",
+      ]
+    attention_head_dim: 32
+    cross_attention_dim: 768
+    

dreamvoice/src/configs/plugin_cross.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/extract_features.py

@@ -0,0 +1,103 @@
+import os
+import torch
+import librosa
+import numpy as np
+import soundfile as sf
+import pandas as pd
+# from feats.hubert_model import get_soft_model, get_hubert_soft_content
+from feats.contentvec_hf import get_content_model, get_content
+# from modules.speaker_encoder.encoder import inference as spk_encoder
+# from pathlib import Path
+from tqdm import tqdm
+from multiprocessing import Process
+import pyworld as pw
+
+
+def resample_save(infolder, audio_path, model,
+                  audio_sr=24000, content_sr=16000, min_length=1.92,
+                  content_resolution=50,
+                  save_path='features'):
+    if os.path.isfile(save_path + '/' + 'audio_24k/' + audio_path) is False:
+        audio, sr = librosa.load(infolder + audio_path, sr=content_sr)
+        final_length = audio.shape[-1] // (content_sr / content_resolution) * (content_sr / content_resolution)
+        # final_length = final_length / content_sr
+
+        length = max(round(min_length*content_sr), round(final_length))
+        assert length % 10 == 0
+        audio = audio[:length]
+        audio_save = np.zeros(length, dtype=audio.dtype)
+        audio_save[:audio.shape[-1]] = audio[:audio.shape[-1]]
+
+        # content = get_hubert_soft_content(model, torch.tensor(audio_save).unsqueeze(0))
+        content = get_content(model, torch.tensor(audio_save).unsqueeze(0))
+        content = content.cpu()
+        os.makedirs(os.path.dirname(save_path + '/' + 'content/' + audio_path), exist_ok=True)
+        torch.save(content, save_path + '/' + 'content/' + audio_path+'.pt')
+        # print(audio_save.shape)
+        # print(content.shape)
+        os.makedirs(os.path.dirname(save_path + '/' + 'audio_16k/' + audio_path), exist_ok=True)
+        sf.write(save_path + '/' + 'audio_16k/' + audio_path, audio_save, int(sr))
+        # print(save_path + '/' + 'audio_16k/' + audio_path)
+
+        audio, sr = librosa.load(infolder + audio_path, sr=audio_sr)
+        length = max(round(min_length*audio_sr), round(final_length/content_sr*audio_sr))
+        assert length % 10 == 0
+        audio = audio[:length]
+        audio_save = np.zeros(length, dtype=audio.dtype)
+        audio_save[:audio.shape[-1]] = audio[:audio.shape[-1]]
+        # print(audio_save.shape)
+        os.makedirs(os.path.dirname(save_path + '/' + 'audio_24k/' + audio_path), exist_ok=True)
+        sf.write(save_path + '/' + 'audio_24k/' + audio_path, audio_save, int(sr))
+
+
+def extract_f0(in_folder, audio_path, save_path):
+    audio, sr = librosa.load(in_folder + audio_path, sr=None)
+    assert sr == 16000
+    if os.path.isfile(save_path + '/' + 'f0/' + audio_path + '.pt') is False:
+        # wav = audio
+        # wav = np.pad(wav, int((1024-320)/2), mode='reflect')
+        # f0_, _, _ = librosa.pyin(wav, frame_length=1024, hop_length=320, center=False, sr=sr,
+        #                         fmin=librosa.note_to_hz('C2'),
+        #                         fmax=librosa.note_to_hz('C6'))
+
+        _f0, t = pw.dio(audio.astype(np.float64), sr, frame_period=320 / sr * 1000)
+        f0 = pw.stonemask(audio.astype(np.float64), _f0, t, sr)[:-1]
+
+        f0 = np.nan_to_num(f0)
+        os.makedirs(os.path.dirname(save_path + '/' + 'f0/' + audio_path), exist_ok=True)
+        # print(save_path + '/' + 'f0/' + audio_path + '.pt')
+        torch.save(torch.tensor(f0), save_path + '/' + 'f0/' + audio_path + '.pt')
+
+
+def chunks(arr, m):
+    result = [[] for i in range(m)]
+    for i in range(len(arr)):
+        result[i%m].append(arr[i])
+    return result
+
+
+def extract_f0_main(in_folder, audio_paths, save_path):
+    for audio_path in tqdm(audio_paths):
+        extract_f0(in_folder, audio_path, save_path)
+
+
+if __name__ == '__main__':
+    df = pd.read_csv('../test_data/vc_meta.csv')
+    # model = get_soft_model('../pre_ckpts/hubert_soft.pt').to('cuda')
+    model = get_content_model().to('cuda')
+    # # spk_encoder.load_model(Path('ckpts/spk_encoder/pretrained.pt'), device="cuda")
+    for i in tqdm(range(len(df))):
+        row = df.iloc[i]
+        in_path = row['path']
+        resample_save('../test_data/', in_path, model, save_path='../features/')
+
+    in_folder = '../features/audio_16k/'
+    audio_files = list(df['path'])
+    save_path = '../features/'
+    cores = 6
+
+    subsets = chunks(audio_files, cores)
+
+    for subset in subsets:
+        t = Process(target=extract_f0_main, args=(in_folder, subset, save_path))
+        t.start()

dreamvoice/src/feats/.ipynb_checkpoints/contentvec-checkpoint.py

@@ -0,0 +1,42 @@
+import torch
+import librosa
+from fairseq import checkpoint_utils
+import torch.nn.functional as F
+
+
+def get_model(vec_path):
+    print("load model(s) from {}".format(vec_path))
+    models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
+        [vec_path],
+        suffix="",
+    )
+    model = models[0]
+    model.eval()
+    return model
+
+
+@torch.no_grad()
+def get_content(hmodel, wav_16k_tensor, device='cuda', layer=12):
+    # print(layer)
+    wav_16k_tensor = wav_16k_tensor.to(device)
+    # so that the output shape will be len(audio//320)
+    wav_16k_tensor = F.pad(wav_16k_tensor, ((400 - 320) // 2, (400 - 320) // 2))
+    feats = wav_16k_tensor
+    padding_mask = torch.BoolTensor(feats.shape).fill_(False)
+    inputs = {
+        "source": feats.to(wav_16k_tensor.device),
+        "padding_mask": padding_mask.to(wav_16k_tensor.device),
+        "output_layer": layer
+    }
+    logits = hmodel.extract_features(**inputs)[0]
+    # feats = hmodel.final_proj(logits[0])
+    return logits
+
+
+if __name__ == '__main__':
+    audio, sr = librosa.load('test.wav', sr=16000)
+    audio = audio[:100*320]
+    model = get_model('../../ckpts/checkpoint_best_legacy_500.pt')
+    model = model.cuda()
+    content = get_content(model, torch.tensor([audio]))
+    print(content)

dreamvoice/src/feats/.ipynb_checkpoints/contentvec_hf-checkpoint.py

@@ -0,0 +1,40 @@
+from transformers import HubertModel
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import librosa
+
+
+class HubertModelWithFinalProj(HubertModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        # The final projection layer is only used for backward compatibility.
+        # Following https://github.com/auspicious3000/contentvec/issues/6
+        # Remove this layer is necessary to achieve the desired outcome.
+        self.final_proj = nn.Linear(config.hidden_size, config.classifier_proj_size)
+
+
+def get_content_model(config='lengyue233/content-vec-best'):
+    model = HubertModelWithFinalProj.from_pretrained(config)
+    model.eval()
+    return model
+
+
+@torch.no_grad()
+def get_content(model, wav_16k_tensor, device='cuda'):
+    # print(layer)
+    wav_16k_tensor = wav_16k_tensor.to(device)
+    # so that the output shape will be len(audio//320)
+    wav_16k_tensor = F.pad(wav_16k_tensor, ((400 - 320) // 2, (400 - 320) // 2))
+    logits = model(wav_16k_tensor)['last_hidden_state']
+    return logits
+
+
+if __name__ == '__main__':
+    model = get_content_model().cuda()
+    audio, sr = librosa.load('test.wav', sr=16000)
+    audio = audio[:100*320]
+    audio = torch.tensor([audio])
+    content = get_content(model, audio, 'cuda')
+    print(content)

dreamvoice/src/feats/.ipynb_checkpoints/hubert_model-checkpoint.py

@@ -0,0 +1,24 @@
+import torch, torchaudio
+from .hubert.hubert import HubertSoft
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+import librosa
+
+
+def get_soft_model(model_path):
+    hubert = HubertSoft()
+    # Load checkpoint (either hubert_soft or hubert_discrete)
+    # hubert = torch.hub.load("bshall/hubert:main", "hubert_soft", trust_repo=True)
+    checkpoint = torch.load(model_path)
+    consume_prefix_in_state_dict_if_present(checkpoint["hubert"], "module.")
+    hubert.load_state_dict(checkpoint["hubert"])
+    hubert.eval()
+    return hubert
+
+
+@torch.no_grad()
+def get_hubert_soft_content(hmodel, wav_16k_tensor, device='cuda'):
+    wav_16k_tensor = wav_16k_tensor.to(device).unsqueeze(1)
+    # print(wav_16k_tensor.shape)
+    units = hmodel.units(wav_16k_tensor)
+    # print(units.shape)
+    return units.cpu()

dreamvoice/src/feats/.ipynb_checkpoints/test-checkpoint.py

@@ -0,0 +1,22 @@
+import torch, torchaudio
+from hubert.hubert import HubertSoft
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+import librosa
+
+
+def get_soft_model(model_path):
+    hubert = HubertSoft()
+    # Load checkpoint (either hubert_soft or hubert_discrete)
+    # hubert = torch.hub.load("bshall/hubert:main", "hubert_soft", trust_repo=True)
+    checkpoint = torch.load(model_path)
+    consume_prefix_in_state_dict_if_present(checkpoint["hubert"], "module.")
+    hubert.load_state_dict(checkpoint["hubert"])
+    hubert.eval()
+    return model
+
+
+@torch.no_grad()
+def get_hubert_soft_content(hmodel, wav_16k_tensor, device='cuda'):
+    wav_16k_tensor = wav_16k_tensor.to(device)
+    units = hmodel.units(wav_16k_tensor)
+    return units.cpu()

dreamvoice/src/feats/contentvec.py

@@ -0,0 +1,42 @@
+import torch
+import librosa
+from fairseq import checkpoint_utils
+import torch.nn.functional as F
+
+
+def get_model(vec_path):
+    print("load model(s) from {}".format(vec_path))
+    models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
+        [vec_path],
+        suffix="",
+    )
+    model = models[0]
+    model.eval()
+    return model
+
+
+@torch.no_grad()
+def get_content(hmodel, wav_16k_tensor, device='cuda', layer=12):
+    # print(layer)
+    wav_16k_tensor = wav_16k_tensor.to(device)
+    # so that the output shape will be len(audio//320)
+    wav_16k_tensor = F.pad(wav_16k_tensor, ((400 - 320) // 2, (400 - 320) // 2))
+    feats = wav_16k_tensor
+    padding_mask = torch.BoolTensor(feats.shape).fill_(False)
+    inputs = {
+        "source": feats.to(wav_16k_tensor.device),
+        "padding_mask": padding_mask.to(wav_16k_tensor.device),
+        "output_layer": layer
+    }
+    logits = hmodel.extract_features(**inputs)[0]
+    # feats = hmodel.final_proj(logits[0])
+    return logits
+
+
+if __name__ == '__main__':
+    audio, sr = librosa.load('test.wav', sr=16000)
+    audio = audio[:100*320]
+    model = get_model('../../ckpts/checkpoint_best_legacy_500.pt')
+    model = model.cuda()
+    content = get_content(model, torch.tensor([audio]))
+    print(content)

dreamvoice/src/feats/contentvec_hf.py

@@ -0,0 +1,40 @@
+from transformers import HubertModel
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import librosa
+
+
+class HubertModelWithFinalProj(HubertModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        # The final projection layer is only used for backward compatibility.
+        # Following https://github.com/auspicious3000/contentvec/issues/6
+        # Remove this layer is necessary to achieve the desired outcome.
+        self.final_proj = nn.Linear(config.hidden_size, config.classifier_proj_size)
+
+
+def get_content_model(config='lengyue233/content-vec-best'):
+    model = HubertModelWithFinalProj.from_pretrained(config)
+    model.eval()
+    return model
+
+
+@torch.no_grad()
+def get_content(model, wav_16k_tensor, device='cuda'):
+    # print(layer)
+    wav_16k_tensor = wav_16k_tensor.to(device)
+    # so that the output shape will be len(audio//320)
+    wav_16k_tensor = F.pad(wav_16k_tensor, ((400 - 320) // 2, (400 - 320) // 2))
+    logits = model(wav_16k_tensor)['last_hidden_state']
+    return logits
+
+
+if __name__ == '__main__':
+    model = get_content_model().cuda()
+    audio, sr = librosa.load('test.wav', sr=16000)
+    audio = audio[:100*320]
+    audio = torch.tensor([audio])
+    content = get_content(model, audio, 'cuda')
+    print(content)

dreamvoice/src/feats/hubert/.gitignore

@@ -0,0 +1,132 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# VSCode project settings
+.vscode
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

dreamvoice/src/feats/hubert/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Benjamin van Niekerk
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.