xinference 1.2.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

xinference/model/llm/vllm/xavier/transfer.py

@@ -23,6 +23,8 @@ from vllm.core.scheduler import Scheduler
 from vllm.utils import TORCH_DTYPE_TO_NUMPY_DTYPE, Device
 from vllm.worker.cache_engine import CacheEngine
 
+from .collective import CollectiveRank
+
 logger = logging.getLogger(__name__)
 
 
@@ -89,7 +91,7 @@ class BufferTransferMixin:
         return TypeMappingGloo[TORCH_DTYPE_TO_NUMPY_DTYPE[input_dtype]]
 
 
-class TransferActor(xo.StatelessActor, BufferTransferMixin):
+class TransferActor(xo.StatelessActor, BufferTransferMixin, CollectiveRank):
     @classmethod
     def default_uid(cls):
         return f"vllm-transfer-actor"
@@ -104,38 +106,21 @@ class TransferActor(xo.StatelessActor, BufferTransferMixin):
         world_addresses: List[str],
     ):
         super().__init__()
-        self._rank = rank
-        self._world_size = world_size
-        self._store_address = store_address
-        self._rank_address = rank_address
-        self._store_port = store_port
-        self._world_addresses = world_addresses
-        self._context = None
+        CollectiveRank.__init__(
+            self,
+            rank,
+            world_size,
+            rank_address,
+            store_address,
+            store_port,
+            world_addresses,
+        )
         self._cache_engine: Optional[List[CacheEngine]] = None
         self._scheduler: Optional[List[Scheduler]] = None
         self._swap_stream = torch.cuda.Stream()
 
     async def __post_create__(self):
-        from xoscar.collective import xoscar_pygloo as xp
-
-        context = xp.rendezvous.Context(self._rank, self._world_size)
-
-        attr = xp.transport.tcp.attr(self._rank_address.split(":")[0])
-        dev = xp.transport.tcp.CreateDevice(attr)
-
-        opt = xp.rendezvous.TCPStoreOptions()
-        opt.port = self._store_port
-        opt.numWorkers = self._world_size
-        opt.isServer = self._rank == 0
-
-        store = xp.rendezvous.TCPStore(self._store_address, opt)
-        store = xp.rendezvous.PrefixStore(str(self._world_size), store)
-
-        context.connectFullMesh(store, dev)
-        self._context = context
-        logger.debug(
-            f"Rank {self._rank} arrives successfully, world addresses: {self._world_addresses}"
-        )
+        self.init_rank()
 
     def setup(
         self,
@@ -153,6 +138,9 @@ class TransferActor(xo.StatelessActor, BufferTransferMixin):
             num_buffer, buffer_shape, buffer_dtype, buffer_device, pin_memory
         )
 
+    async def __pre_destroy__(self):
+        self._context.closeConnections()
+
     def _get_cache_engine(self, virtual_engine: int) -> CacheEngine:
         return self._cache_engine[virtual_engine]  # type: ignore
 
@@ -281,18 +269,51 @@ class TransferActor(xo.StatelessActor, BufferTransferMixin):
             self.free_buffer_index(cpu_buf_index)
 
     async def recv(
-        self, virtual_engine: int, from_address: str, src_to_dst: Dict[int, int]
+        self, virtual_engine: int, from_rank: int, src_to_dst: Dict[int, int]
     ):
         """
         This is the external entry point for the call.
         The transfer logic is as follows:
         the receiver requests the sender to send the data directly to itself in a point-to-point manner.
         """
-        rank = self._world_addresses.index(from_address)
+        from_address = self._world_addresses[from_rank]
         sender_ref = await xo.actor_ref(
-            address=from_address, uid=f"{TransferActor.default_uid()}-{rank}"
+            address=from_address, uid=f"{TransferActor.default_uid()}-{from_rank}"
         )
         await asyncio.gather(
             sender_ref.do_send(virtual_engine, self._rank, src_to_dst),
-            self.do_recv(virtual_engine, rank, src_to_dst),
+            self.do_recv(virtual_engine, from_rank, src_to_dst),
         )
+
+
+class Rank0TransferActor(xo.StatelessActor, CollectiveRank):
+    """
+    The Rank 0 transfer actor is only used for constructing the collective communication world,
+    so it only needs to inherit the `CollectiveWorld` class.
+    """
+
+    @classmethod
+    def default_uid(cls):
+        return f"vllm-transfer-actor"
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        rank_address: str,
+        store_address: str,
+        store_port: int,
+        world_addresses: List[str],
+    ):
+        CollectiveRank.__init__(
+            self,
+            rank,
+            world_size,
+            rank_address,
+            store_address,
+            store_port,
+            world_addresses,
+        )
+
+    async def __post_create__(self):
+        self.init_rank()

xinference/thirdparty/melo/api.py

@@ -0,0 +1,135 @@
+import os
+import re
+import json
+import torch
+import librosa
+import soundfile
+import torchaudio
+import numpy as np
+import torch.nn as nn
+from tqdm import tqdm
+import torch
+
+from . import utils
+from . import commons
+from .models import SynthesizerTrn
+from .split_utils import split_sentence
+from .mel_processing import spectrogram_torch, spectrogram_torch_conv
+from .download_utils import load_or_download_config, load_or_download_model
+
+class TTS(nn.Module):
+    def __init__(self, 
+                language,
+                device='auto',
+                use_hf=True,
+                config_path=None,
+                ckpt_path=None):
+        super().__init__()
+        if device == 'auto':
+            device = 'cpu'
+            if torch.cuda.is_available(): device = 'cuda'
+            if torch.backends.mps.is_available(): device = 'mps'
+        if 'cuda' in device:
+            assert torch.cuda.is_available()
+
+        # config_path = 
+        hps = load_or_download_config(language, use_hf=use_hf, config_path=config_path)
+
+        num_languages = hps.num_languages
+        num_tones = hps.num_tones
+        symbols = hps.symbols
+
+        model = SynthesizerTrn(
+            len(symbols),
+            hps.data.filter_length // 2 + 1,
+            hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            num_tones=num_tones,
+            num_languages=num_languages,
+            **hps.model,
+        ).to(device)
+
+        model.eval()
+        self.model = model
+        self.symbol_to_id = {s: i for i, s in enumerate(symbols)}
+        self.hps = hps
+        self.device = device
+    
+        # load state_dict
+        checkpoint_dict = load_or_download_model(language, device, use_hf=use_hf, ckpt_path=ckpt_path)
+        self.model.load_state_dict(checkpoint_dict['model'], strict=True)
+        
+        language = language.split('_')[0]
+        self.language = 'ZH_MIX_EN' if language == 'ZH' else language # we support a ZH_MIX_EN model
+
+    @staticmethod
+    def audio_numpy_concat(segment_data_list, sr, speed=1.):
+        audio_segments = []
+        for segment_data in segment_data_list:
+            audio_segments += segment_data.reshape(-1).tolist()
+            audio_segments += [0] * int((sr * 0.05) / speed)
+        audio_segments = np.array(audio_segments).astype(np.float32)
+        return audio_segments
+
+    @staticmethod
+    def split_sentences_into_pieces(text, language, quiet=False):
+        texts = split_sentence(text, language_str=language)
+        if not quiet:
+            print(" > Text split to sentences.")
+            print('\n'.join(texts))
+            print(" > ===========================")
+        return texts
+
+    def tts_to_file(self, text, speaker_id, output_path=None, sdp_ratio=0.2, noise_scale=0.6, noise_scale_w=0.8, speed=1.0, pbar=None, format=None, position=None, quiet=False,):
+        language = self.language
+        texts = self.split_sentences_into_pieces(text, language, quiet)
+        audio_list = []
+        if pbar:
+            tx = pbar(texts)
+        else:
+            if position:
+                tx = tqdm(texts, position=position)
+            elif quiet:
+                tx = texts
+            else:
+                tx = tqdm(texts)
+        for t in tx:
+            if language in ['EN', 'ZH_MIX_EN']:
+                t = re.sub(r'([a-z])([A-Z])', r'\1 \2', t)
+            device = self.device
+            bert, ja_bert, phones, tones, lang_ids = utils.get_text_for_tts_infer(t, language, self.hps, device, self.symbol_to_id)
+            with torch.no_grad():
+                x_tst = phones.to(device).unsqueeze(0)
+                tones = tones.to(device).unsqueeze(0)
+                lang_ids = lang_ids.to(device).unsqueeze(0)
+                bert = bert.to(device).unsqueeze(0)
+                ja_bert = ja_bert.to(device).unsqueeze(0)
+                x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
+                del phones
+                speakers = torch.LongTensor([speaker_id]).to(device)
+                audio = self.model.infer(
+                        x_tst,
+                        x_tst_lengths,
+                        speakers,
+                        tones,
+                        lang_ids,
+                        bert,
+                        ja_bert,
+                        sdp_ratio=sdp_ratio,
+                        noise_scale=noise_scale,
+                        noise_scale_w=noise_scale_w,
+                        length_scale=1. / speed,
+                    )[0][0, 0].data.cpu().float().numpy()
+                del x_tst, tones, lang_ids, bert, ja_bert, x_tst_lengths, speakers
+                # 
+            audio_list.append(audio)
+        torch.cuda.empty_cache()
+        audio = self.audio_numpy_concat(audio_list, sr=self.hps.data.sampling_rate, speed=speed)
+
+        if output_path is None:
+            return audio
+        else:
+            if format:
+                soundfile.write(output_path, audio, self.hps.data.sampling_rate, format=format)
+            else:
+                soundfile.write(output_path, audio, self.hps.data.sampling_rate)

xinference/thirdparty/melo/app.py

@@ -0,0 +1,61 @@
+# WebUI by mrfakename <X @realmrfakename / HF @mrfakename>
+# Demo also available on HF Spaces: https://huggingface.co/spaces/mrfakename/MeloTTS
+import gradio as gr
+import os, torch, io
+# os.system('python -m unidic download')
+print("Make sure you've downloaded unidic (python -m unidic download) for this WebUI to work.")
+from melo.api import TTS
+speed = 1.0
+import tempfile
+import click
+device = 'auto'
+models = {
+    'EN': TTS(language='EN', device=device),
+    'ES': TTS(language='ES', device=device),
+    'FR': TTS(language='FR', device=device),
+    'ZH': TTS(language='ZH', device=device),
+    'JP': TTS(language='JP', device=device),
+    'KR': TTS(language='KR', device=device),
+}
+speaker_ids = models['EN'].hps.data.spk2id
+
+default_text_dict = {
+    'EN': 'The field of text-to-speech has seen rapid development recently.',
+    'ES': 'El campo de la conversión de texto a voz ha experimentado un rápido desarrollo recientemente.',
+    'FR': 'Le domaine de la synthèse vocale a connu un développement rapide récemment',
+    'ZH': 'text-to-speech 领域近年来发展迅速',
+    'JP': 'テキスト読み上げの分野は最近急速な発展を遂げています',
+    'KR': '최근 텍스트 음성 변환 분야가 급속도로 발전하고 있습니다.',    
+}
+    
+def synthesize(speaker, text, speed, language, progress=gr.Progress()):
+    bio = io.BytesIO()
+    models[language].tts_to_file(text, models[language].hps.data.spk2id[speaker], bio, speed=speed, pbar=progress.tqdm, format='wav')
+    return bio.getvalue()
+def load_speakers(language, text):
+    if text in list(default_text_dict.values()):
+        newtext = default_text_dict[language]
+    else:
+        newtext = text
+    return gr.update(value=list(models[language].hps.data.spk2id.keys())[0], choices=list(models[language].hps.data.spk2id.keys())), newtext
+with gr.Blocks() as demo:
+    gr.Markdown('# MeloTTS WebUI\n\nA WebUI for MeloTTS.')
+    with gr.Group():
+        speaker = gr.Dropdown(speaker_ids.keys(), interactive=True, value='EN-US', label='Speaker')
+        language = gr.Radio(['EN', 'ES', 'FR', 'ZH', 'JP', 'KR'], label='Language', value='EN')
+        speed = gr.Slider(label='Speed', minimum=0.1, maximum=10.0, value=1.0, interactive=True, step=0.1)
+        text = gr.Textbox(label="Text to speak", value=default_text_dict['EN'])
+        language.input(load_speakers, inputs=[language, text], outputs=[speaker, text])
+    btn = gr.Button('Synthesize', variant='primary')
+    aud = gr.Audio(interactive=False)
+    btn.click(synthesize, inputs=[speaker, text, speed, language], outputs=[aud])
+    gr.Markdown('WebUI by [mrfakename](https://twitter.com/realmrfakename).')
+@click.command()
+@click.option('--share', '-s', is_flag=True, show_default=True, default=False, help="Expose a publicly-accessible shared Gradio link usable by anyone with the link. Only share the link with people you trust.")
+@click.option('--host', '-h', default=None)
+@click.option('--port', '-p', type=int, default=None)
+def main(share, host, port):
+    demo.queue(api_open=False).launch(show_api=False, share=share, server_name=host, server_port=port)
+
+if __name__ == "__main__":
+    main()