Delete CHATTS/core.py

CHATTS/core.py

@@ -1,149 +0,0 @@
-
-import os
-import logging
-from omegaconf import OmegaConf
-
-import torch
-from vocos import Vocos
-from .model.dvae import DVAE
-from .model.gpt import GPT_warpper
-from .utils.gpu_utils import select_device
-from .utils.io_utils import get_latest_modified_file
-from .infer.api import refine_text, infer_code
-
-from huggingface_hub import snapshot_download
-
-logging.basicConfig(level = logging.INFO)
-
-
-class Chat:
-    def __init__(self, ):
-        self.pretrain_models = {}
-        self.logger = logging.getLogger(__name__)
-        
-    def check_model(self, level = logging.INFO, use_decoder = False):
-        not_finish = False
-        check_list = ['vocos', 'gpt', 'tokenizer']
-        
-        if use_decoder:
-            check_list.append('decoder')
-        else:
-            check_list.append('dvae')
-            
-        for module in check_list:
-            if module not in self.pretrain_models:
-                self.logger.log(logging.WARNING, f'{module} not initialized.')
-                not_finish = True
-                
-        if not not_finish:
-            self.logger.log(level, f'All initialized.')
-            
-        return not not_finish
-        
-    def load_models(self, source='huggingface', force_redownload=False, local_path='<LOCAL_PATH>'):
-        if source == 'huggingface':
-            hf_home = os.getenv('HF_HOME', os.path.expanduser("~/.cache/huggingface"))
-            try:
-                download_path = get_latest_modified_file(os.path.join(hf_home, 'hub/models--2Noise--ChatTTS/snapshots'))
-            except:
-                download_path = None
-            if download_path is None or force_redownload: 
-                self.logger.log(logging.INFO, f'Download from HF: https://huggingface.co/2Noise/ChatTTS')
-                download_path = snapshot_download(repo_id="2Noise/ChatTTS", allow_patterns=["*.pt", "*.yaml"])
-            else:
-                self.logger.log(logging.INFO, f'Load from cache: {download_path}')
-            self._load(**{k: os.path.join(download_path, v) for k, v in OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()})
-        elif source == 'local':
-            self.logger.log(logging.INFO, f'Load from local: {local_path}')
-            self._load(**{k: os.path.join(local_path, v) for k, v in OmegaConf.load(os.path.join(local_path, 'config', 'path.yaml')).items()})
-        
-    def _load(
-        self, 
-        vocos_config_path: str = None, 
-        vocos_ckpt_path: str = None,
-        dvae_config_path: str = None,
-        dvae_ckpt_path: str = None,
-        gpt_config_path: str = None,
-        gpt_ckpt_path: str = None,
-        decoder_config_path: str = None,
-        decoder_ckpt_path: str = None,
-        tokenizer_path: str = None,
-        device: str = None
-    ):
-        if not device:
-            device = select_device(4096)
-            self.logger.log(logging.INFO, f'use {device}')
-            
-        if vocos_config_path:
-            vocos = Vocos.from_hparams(vocos_config_path).to(device).eval()
-            assert vocos_ckpt_path, 'vocos_ckpt_path should not be None'
-            vocos.load_state_dict(torch.load(vocos_ckpt_path))
-            self.pretrain_models['vocos'] = vocos
-            self.logger.log(logging.INFO, 'vocos loaded.')
-        
-        if dvae_config_path:
-            cfg = OmegaConf.load(dvae_config_path)
-            dvae = DVAE(**cfg).to(device).eval()
-            assert dvae_ckpt_path, 'dvae_ckpt_path should not be None'
-            dvae.load_state_dict(torch.load(dvae_ckpt_path, map_location='cpu'))
-            self.pretrain_models['dvae'] = dvae
-            self.logger.log(logging.INFO, 'dvae loaded.')
-            
-        if gpt_config_path:
-            cfg = OmegaConf.load(gpt_config_path)
-            gpt = GPT_warpper(**cfg).to(device).eval()
-            assert gpt_ckpt_path, 'gpt_ckpt_path should not be None'
-            gpt.load_state_dict(torch.load(gpt_ckpt_path, map_location='cpu'))
-            self.pretrain_models['gpt'] = gpt
-            self.logger.log(logging.INFO, 'gpt loaded.')
-            
-        if decoder_config_path:
-            cfg = OmegaConf.load(decoder_config_path)
-            decoder = DVAE(**cfg).to(device).eval()
-            assert decoder_ckpt_path, 'decoder_ckpt_path should not be None'
-            decoder.load_state_dict(torch.load(decoder_ckpt_path, map_location='cpu'))
-            self.pretrain_models['decoder'] = decoder
-            self.logger.log(logging.INFO, 'decoder loaded.')
-        
-        if tokenizer_path:
-            tokenizer = torch.load(tokenizer_path, map_location='cpu')
-            tokenizer.padding_side = 'left'
-            self.pretrain_models['tokenizer'] = tokenizer
-            self.logger.log(logging.INFO, 'tokenizer loaded.')
-            
-        self.check_model()
-    
-    def infer(
-        self, 
-        text, 
-        skip_refine_text=False, 
-        refine_text_only=False, 
-        params_refine_text={}, 
-        params_infer_code={}, 
-        use_decoder=False
-    ):
-        
-        assert self.check_model(use_decoder=use_decoder)
-        
-        if not skip_refine_text:
-            text_tokens = refine_text(self.pretrain_models, text, **params_refine_text)['ids']
-            text_tokens = [i[i < self.pretrain_models['tokenizer'].convert_tokens_to_ids('[break_0]')] for i in text_tokens]
-            text = self.pretrain_models['tokenizer'].batch_decode(text_tokens)
-            if refine_text_only:
-                return text
-            
-        text = [params_infer_code.get('prompt', '') + i for i in text]
-        params_infer_code.pop('prompt', '')
-        result = infer_code(self.pretrain_models, text, **params_infer_code, return_hidden=use_decoder)
-        
-        if use_decoder:
-            mel_spec = [self.pretrain_models['decoder'](i[None].permute(0,2,1)) for i in result['hiddens']]
-        else:
-            mel_spec = [self.pretrain_models['dvae'](i[None].permute(0,2,1)) for i in result['ids']]
-            
-        wav = [self.pretrain_models['vocos'].decode(i).cpu().numpy() for i in mel_spec]
-        
-        return wav
-        
-        
-