Coverage for wrapper/xtts/wrapper_xtts.py: 75%

1import logging 

2 

3import torch 

4from torch import inference_mode 

5 

6from typing_extensions import override # Python 3.11 

7from typing import Optional 

8from typing import Dict 

9from typing import Any 

10from typing import Union 

11 

12from wrapper_model import ModelGeneration 

13 

14import numpy as np 

15 

16from TTS.tts.configs.xtts_config import XttsConfig 

17from TTS.tts.models.xtts import Xtts 

18 

19 

20class XTTSGeneration(ModelGeneration): 

21 def __init__(self) -> None: 

22 super().__init__("xtts") 

23 

24 # Model components 

25 self.xtts_config: Optional[XttsConfig] = None 

26 self.xtts: Optional[Xtts] = None 

27 

28 def __del__(self) -> None: 

29 if self.xtts is not None: 

30 self.xtts = None 

31 

32 def init_parallelism(self) -> None: 

33 self.load_timer.start("torch_dist") 

34 # No real parallelism as it runs with a single GPU or CPU 

35 if torch.cuda.is_available(): 

36 self.rank = 0 

37 self.local_rank = 0 

38 self.world_size = 1 

39 self.device_id: Union[int, str] = self.local_rank 

40 self.device = torch.device(f"cuda:{self.device_id}") 

41 torch.cuda.set_device(self.local_rank) 

42 else: 

43 self.device_id = "cpu" 

44 self.device = torch.device(self.device_id) 

45 self.load_timer.end("torch_dist") 

46 

47 def load_model(self) -> None: 

48 self.load_timer.start("xtts") 

49 self.xtts_config = XttsConfig() 

50 # TODO fix this with the DockerFile 

51 self.xtts_config.load_json("coqui/XTTS-v2/config.json") 

52 self.xtts = Xtts.init_from_config(self.xtts_config) 

53 self.xtts.load_checkpoint( 

54 self.xtts_config, 

55 checkpoint_dir="coqui/XTTS-v2", 

56 eval=True) 

57 self.xtts.cuda() 

58 self.load_timer.end("xtts") 

59 

60 def init_model_parallelism(self) -> None: 

61 if self.world_size > 1: 

62 logging.warning("XTTS does not support distributed parallelism.") 

63 

64 def model_compile(self) -> None: 

65 if not self.torch_compile: 

66 return 

67 self.load_timer.start("compile") 

68 self.xtts = torch.compile( 

69 self.xtts, 

70 mode="reduce-overhead") 

71 self.load_timer.end("compile") 

72 

73 async def get_rest_args( 

74 self, 

75 data_json: Dict[str, Union[str, int, float]] 

76 ) -> Dict[str, Any]: 

77 if data_json is None: 

78 raise ValueError("Missing JSON body") 

79 text = data_json.get("text", None) 

80 if text is None: 

81 raise ValueError("Missing 'text' parameter") 

82 return { 

83 "task": self.model_name, 

84 "args": { 

85 "text": text, 

86 } 

87 } 

88 

89 @inference_mode() 

90 async def warmup(self) -> None: 

91 logging.info("Warmup for XTTS generation") 

92 await self.generate(text="Warmup") 

93 

94 def _assert_model_init(self) -> None: 

95 super()._assert_model_init() 

96 if not self.xtts: 

97 raise ValueError("XTTS not loaded.") 

98 

99 @override 

100 @inference_mode() 

101 async def generate( 

102 self, 

103 text: str, 

104 job_id: Optional[str] = None, 

105 ) -> np.ndarray: 

106 gen_timer = self._new_gen_timer(job_id) 

107 

108 self._assert_model_init() 

109 

110 self.running = True # We can run in parallel but good to know if we are running 

111 

112 try: 

113 # Clean up text 

114 text = text.replace("*", "") 

115 

116 if self.xtts is None: 

117 raise RuntimeError("Model not loaded") 

118 outputs = self.xtts.synthesize( 

119 text, 

120 self.xtts_config, 

121 speaker_wav="tests/data/ljspeech/wavs/LJ001-0001.wav", 

122 gpt_cond_len=3, 

123 language="en", 

124 ) 

125 wav_np = outputs["wav"] 

126 return wav_np 

127 finally: 

128 self.running = False 

129 gen_timer.end("total") 

130 

131 def get_health(self) -> Dict[str, Any]: 

132 ret = super().get_health() 

133 if torch.cuda.is_available(): 

134 ret["gpu"] = torch.cuda.get_device_name(0) 

135 return ret