fix: mps memory leakage by moving gpt&emb to cpu (#389)

see more at
#261
pytorch/pytorch#91368
https://stackoverflow.com/questions/74902903/pytorch-memory-leak-reference-cycle-in-for-loop

Author: fumiama

ChatTTS/core.py

@@ -5,6 +5,7 @@
 from functools import partial
 from typing import Literal
 import tempfile
+from typing import Optional
 
 import torch
 from omegaconf import OmegaConf
@@ -15,7 +16,7 @@
 from .model.gpt import GPT_warpper
 from .utils.gpu_utils import select_device
 from .utils.infer_utils import count_invalid_characters, detect_language, apply_character_map, apply_half2full_map, HomophonesReplacer
-from .utils.io_utils import get_latest_modified_file
+from .utils.io import get_latest_modified_file, del_all
 from .infer.api import refine_text, infer_code
 from .utils.download import check_all_assets, download_all_assets
 
@@ -91,17 +92,18 @@ def _load(
         decoder_config_path: str = None,
         decoder_ckpt_path: str = None,
         tokenizer_path: str = None,
-        device: str = None,
+        device: Optional[torch.device] = None,
         compile: bool = True,
     ):
-        if not device:
-            device = select_device(4095)
+        if device is None:
+            device = select_device(4096)
             self.logger.log(logging.INFO, f'use {device}')
-            
+        self.device = device
+
         if vocos_config_path:
             vocos = Vocos.from_hparams(vocos_config_path).to(
                 # vocos on mps will crash, use cpu fallback
-                "cpu" if torch.backends.mps.is_available() else device
+                "cpu" if "mps" in str(device) else device
             ).eval()
             assert vocos_ckpt_path, 'vocos_ckpt_path should not be None'
             vocos.load_state_dict(torch.load(vocos_ckpt_path))
@@ -118,7 +120,7 @@ def _load(
 
         if gpt_config_path:
             cfg = OmegaConf.load(gpt_config_path)
-            gpt = GPT_warpper(**cfg).to(device).eval()
+            gpt = GPT_warpper(**cfg, device=device).eval()
             assert gpt_ckpt_path, 'gpt_ckpt_path should not be None'
             gpt.load_state_dict(torch.load(gpt_ckpt_path))
             if compile and 'cuda' in str(device):
@@ -188,6 +190,7 @@ def _infer(
             text_tokens = refine_text(
                 self.pretrain_models,
                 text,
+                device=self.device,
                 **params_refine_text,
             )['ids']
             text_tokens = [i[i < self.pretrain_models['tokenizer'].convert_tokens_to_ids('[break_0]')] for i in text_tokens]
@@ -198,16 +201,28 @@ def _infer(
 
         text = [params_infer_code.get('prompt', '') + i for i in text]
         params_infer_code.pop('prompt', '')
-        result_gen = infer_code(self.pretrain_models, text, **params_infer_code, return_hidden=use_decoder, stream=stream)
+        result_gen = infer_code(
+            self.pretrain_models,
+            text,
+            device=self.device,
+            **params_infer_code,
+            return_hidden=use_decoder,
+            stream=stream,
+        )
         if use_decoder:
             field = 'hiddens'
             docoder_name = 'decoder'
         else:
             field = 'ids'
             docoder_name = 'dvae'
-        vocos_decode = lambda spec: [self.pretrain_models['vocos'].decode(
-                    i.cpu() if torch.backends.mps.is_available() else i
-                ).cpu().numpy() for i in spec]
+        if "mps" in str(self.device):
+            vocos_decode = lambda spec: [self.pretrain_models['vocos'].decode(
+                i.cpu()
+            ).cpu().numpy() for i in spec]
+        else:
+            vocos_decode = lambda spec: [self.pretrain_models['vocos'].decode(
+                i
+            ).cpu().numpy() for i in spec]
         if stream:
 
             length = 0
@@ -221,13 +236,20 @@ def _infer(
                 if not len(chunk_data):
                     continue
                 self.logger.debug(f'new hidden {len(chunk_data)=}')
-                mel_spec = [self.pretrain_models[docoder_name](i[None].permute(0,2,1)) for i in [chunk_data]]
+                mel_spec = [self.pretrain_models[docoder_name](i[None].permute(0,2,1).to(self.device)) for i in [chunk_data]]
+                del_all(result)
+                del chunk_data
                 wav = vocos_decode(mel_spec)
+                del_all(mel_spec)
                 self.logger.debug(f'yield wav chunk {len(wav[0])=} {len(wav[0][0])=}')
                 yield wav
             return
-        mel_spec = [self.pretrain_models[docoder_name](i[None].permute(0,2,1)) for i in next(result_gen)[field]]
-        yield vocos_decode(mel_spec)
+        result = next(result_gen)
+        mel_spec = [self.pretrain_models[docoder_name](i[None].permute(0,2,1).to(self.device)) for i in result[field]]
+        del_all(result)
+        wav = vocos_decode(mel_spec)
+        del_all(mel_spec)
+        yield wav
 
     def infer(
         self, 

ChatTTS/infer/api.py

@@ -2,7 +2,10 @@
 import torch
 import torch.nn.functional as F
 from transformers.generation import TopKLogitsWarper, TopPLogitsWarper
+
 from ..utils.infer_utils import CustomRepetitionPenaltyLogitsProcessorRepeat
+from ..utils.io import del_all
+from ..model.gpt import GPT_warpper
 
 def infer_code(
     models,
@@ -14,39 +17,42 @@ def infer_code(
     repetition_penalty = 1.05,
     max_new_token = 2048,
     stream=False,
+    device="cpu",
     **kwargs
 ):
-    
-    device = next(models['gpt'].parameters()).device
-    
+
+    gpt: GPT_warpper = models['gpt']
+
     if not isinstance(text, list): 
         text = [text]
 
     if not isinstance(temperature, list):
-        temperature = [temperature] * models['gpt'].num_vq
+        temperature = [temperature] * gpt.num_vq
 
     if spk_emb is not None:
         text = [f'[Stts][spk_emb]{i}[Ptts]' for i in text] 
     else:
         text = [f'[Stts][empty_spk]{i}[Ptts]' for i in text]
 
-    text_token = models['tokenizer'](text, return_tensors='pt', add_special_tokens=False, padding=True).to(device)
-    input_ids = text_token['input_ids'][...,None].expand(-1, -1, models['gpt'].num_vq)
-    text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=device)
-    
-    inputs = {
-        'input_ids': input_ids,
-        'text_mask': text_mask,
-        'attention_mask': text_token['attention_mask'],
-    }
+    text_token_tmp = models['tokenizer'](text, return_tensors='pt', add_special_tokens=False, padding=True)
+    text_token = text_token_tmp.to(device)
+    del text_token_tmp
+    input_ids = text_token['input_ids'][...,None].expand(-1, -1, gpt.num_vq).to(gpt.device_gpt)
+    text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=gpt.device_gpt)
+
+    emb = gpt.get_emb(
+        input_ids=input_ids,
+        text_mask=text_mask,
+    )
+    del text_mask
 
-    emb = models['gpt'].get_emb(**inputs)
     if spk_emb is not None:
-        emb[inputs['input_ids'][..., 0] == models['tokenizer'].convert_tokens_to_ids('[spk_emb]')] = \
-            F.normalize(spk_emb.to(device).to(emb.dtype)[None].expand(len(text), -1), p=2.0, dim=1, eps=1e-12)  
-    
-    num_code = models['gpt'].emb_code[0].num_embeddings - 1
-    
+        n = F.normalize(spk_emb.to(emb.dtype)[None].expand(len(text), -1), p=2.0, dim=1, eps=1e-12).to(gpt.device_gpt)
+        emb[input_ids[..., 0] == models['tokenizer'].convert_tokens_to_ids('[spk_emb]')] = n
+        del n
+
+    num_code = int(gpt.emb_code[0].num_embeddings - 1)
+
     LogitsWarpers = []
     if top_P is not None:
         LogitsWarpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
@@ -58,10 +64,10 @@ def infer_code(
         LogitsProcessors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(\
             repetition_penalty, num_code, 16))
 
-    result = models['gpt'].generate(
-        emb, inputs['input_ids'], 
+    result = gpt.generate(
+        emb, input_ids, 
         temperature = torch.tensor(temperature, device=device), 
-        attention_mask = inputs['attention_mask'],
+        attention_mask = text_token['attention_mask'],
         LogitsWarpers = LogitsWarpers,
         LogitsProcessors = LogitsProcessors,
         eos_token = num_code, 
@@ -71,6 +77,11 @@ def infer_code(
         **kwargs
     )
 
+    del_all(text_token)
+    del emb, text_token, input_ids
+    del_all(LogitsWarpers)
+    del_all(LogitsProcessors)
+
     return result
 
 
@@ -83,11 +94,12 @@ def refine_text(
     repetition_penalty = 1.0,
     max_new_token = 384,
     prompt = '',
+    device="cpu",
     **kwargs
 ):
-    
-    device = next(models['gpt'].parameters()).device
-    
+
+    gpt: GPT_warpper = models['gpt']
+
     if not isinstance(text, list): 
         text = [text]
 
@@ -97,11 +109,7 @@ def refine_text(
     text_token = models['tokenizer'](text, return_tensors='pt', add_special_tokens=False, padding=True).to(device)
     text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=device)
 
-    inputs = {
-        'input_ids': text_token['input_ids'][...,None].expand(-1, -1, models['gpt'].num_vq),
-        'text_mask': text_mask,
-        'attention_mask': text_token['attention_mask'],
-    }
+    input_ids = text_token['input_ids'][...,None].expand(-1, -1, gpt.num_vq)
 
     LogitsWarpers = []
     if top_P is not None:
@@ -112,11 +120,17 @@ def refine_text(
     LogitsProcessors = []
     if repetition_penalty is not None and repetition_penalty != 1:
         LogitsProcessors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(repetition_penalty, len(models['tokenizer']), 16))
-    
-    result = models['gpt'].generate(
-        models['gpt'].get_emb(**inputs), inputs['input_ids'], 
+
+    emb = gpt.get_emb(
+        input_ids=input_ids,
+        text_mask=text_mask,
+    )
+    del text_mask
+
+    result = gpt.generate(
+        emb, input_ids, 
         temperature = torch.tensor([temperature,], device=device), 
-        attention_mask = inputs['attention_mask'],
+        attention_mask = text_token['attention_mask'],
         LogitsWarpers = LogitsWarpers,
         LogitsProcessors = LogitsProcessors,
         eos_token = torch.tensor(models['tokenizer'].convert_tokens_to_ids('[Ebreak]'), device=device)[None], 
@@ -125,4 +139,10 @@ def refine_text(
         stream = False,
         **kwargs
     )
+
+    del_all(text_token)
+    del emb, text_token, input_ids
+    del_all(LogitsWarpers)
+    del_all(LogitsProcessors)
+
     return next(result)