xinference 0.13.0__py3-none-any.whl → 0.13.1__py3-none-any.whl

xinference/thirdparty/ChatTTS/experimental/llm.py

@@ -1,40 +0,0 @@
-
-from openai import OpenAI
- 
-prompt_dict = {
-    'kimi': [ {"role": "system", "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。"},
-              {"role": "user", "content": "你好，请注意你现在生成的文字要按照人日常生活的口吻，你的回复将会后续用TTS模型转为语音，并且请把回答控制在100字以内。并且标点符号仅包含逗号和句号，将数字等转为文字回答。"},
-              {"role": "assistant", "content": "好的，我现在生成的文字将按照人日常生活的口吻， 并且我会把回答控制在一百字以内, 标点符号仅包含逗号和句号，将阿拉伯数字等转为中文文字回答。下面请开始对话。"},],
-    'deepseek': [
-        {"role": "system", "content": "You are a helpful assistant"},
-        {"role": "user", "content": "你好，请注意你现在生成的文字要按照人日常生活的口吻，你的回复将会后续用TTS模型转为语音，并且请把回答控制在100字以内。并且标点符号仅包含逗号和句号，将数字等转为文字回答。"},
-        {"role": "assistant", "content": "好的，我现在生成的文字将按照人日常生活的口吻， 并且我会把回答控制在一百字以内, 标点符号仅包含逗号和句号，将阿拉伯数字等转为中文文字回答。下面请开始对话。"},],
-    'deepseek_TN': [
-        {"role": "system", "content": "You are a helpful assistant"},
-        {"role": "user", "content": "你好，现在我们在处理TTS的文本输入，下面将会给你输入一段文本，请你将其中的阿拉伯数字等等转为文字表达，并且输出的文本里仅包含逗号和句号这两个标点符号"},
-        {"role": "assistant", "content": "好的，我现在对TTS的文本输入进行处理。这一般叫做text normalization。下面请输入"},
-        {"role": "user", "content": "We paid $123 for this desk."},
-        {"role": "assistant", "content": "We paid one hundred and twenty three dollars for this desk."},
-        {"role": "user", "content": "详询请拨打010-724654"},
-        {"role": "assistant", "content": "详询请拨打零幺零，七二四六五四"},
-        {"role": "user", "content": "罗森宣布将于7月24日退市，在华门店超6000家！"},
-        {"role": "assistant", "content": "罗森宣布将于七月二十四日退市，在华门店超过六千家。"},
-        ],
-}          
-                
-class llm_api:
-    def __init__(self, api_key, base_url, model):
-        self.client =  OpenAI(
-            api_key = api_key,
-            base_url = base_url,
-        )
-        self.model = model
-    def call(self, user_question, temperature = 0.3, prompt_version='kimi', **kwargs):
-    
-        completion = self.client.chat.completions.create(
-            model = self.model,
-            messages = prompt_dict[prompt_version]+[{"role": "user", "content": user_question},],
-            temperature = temperature,
-            **kwargs
-        )
-        return completion.choices[0].message.content

xinference/thirdparty/ChatTTS/infer/api.py

@@ -1,125 +0,0 @@
-
-import torch
-import torch.nn.functional as F
-from transformers.generation import TopKLogitsWarper, TopPLogitsWarper
-from ..utils.infer_utils import CustomRepetitionPenaltyLogitsProcessorRepeat
-
-def infer_code(
-    models,
-    text, 
-    spk_emb = None,
-    top_P = 0.7, 
-    top_K = 20, 
-    temperature = 0.3, 
-    repetition_penalty = 1.05,
-    max_new_token = 2048,
-    **kwargs
-):
-    
-    device = next(models['gpt'].parameters()).device
-    
-    if not isinstance(text, list): 
-        text = [text]
-        
-    if not isinstance(temperature, list):
-        temperature = [temperature] * models['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'],
-    }
-
-    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
-    
-    LogitsWarpers = []
-    if top_P is not None:
-        LogitsWarpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
-    if top_K is not None:
-        LogitsWarpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
-        
-    LogitsProcessors = []
-    if repetition_penalty is not None and repetition_penalty != 1:
-        LogitsProcessors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(\
-            repetition_penalty, num_code, 16))
-    
-    result = models['gpt'].generate(
-        emb, inputs['input_ids'], 
-        temperature = torch.tensor(temperature, device=device), 
-        attention_mask = inputs['attention_mask'],
-        LogitsWarpers = LogitsWarpers,
-        LogitsProcessors = LogitsProcessors,
-        eos_token = num_code, 
-        max_new_token = max_new_token, 
-        infer_text = False,
-        **kwargs
-    )
-    
-    return result
-
-
-def refine_text(
-    models, 
-    text,
-    top_P = 0.7, 
-    top_K = 20, 
-    temperature = 0.7, 
-    repetition_penalty = 1.0,
-    max_new_token = 384,
-    prompt = '',
-    **kwargs
-):
-    
-    device = next(models['gpt'].parameters()).device
-    
-    if not isinstance(text, list): 
-        text = [text]
-    
-    assert len(text), 'text should not be empty'
-
-    text = [f"[Sbreak]{i}[Pbreak]{prompt}" for i in 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'],
-    }
-    
-    LogitsWarpers = []
-    if top_P is not None:
-        LogitsWarpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
-    if top_K is not None:
-        LogitsWarpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
-        
-    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'], 
-        temperature = torch.tensor([temperature,], device=device), 
-        attention_mask = inputs['attention_mask'],
-        LogitsWarpers = LogitsWarpers,
-        LogitsProcessors = LogitsProcessors,
-        eos_token = torch.tensor(models['tokenizer'].convert_tokens_to_ids('[Ebreak]'), device=device)[None], 
-        max_new_token = max_new_token, 
-        infer_text = True,
-        **kwargs
-    )
-    return result

xinference/thirdparty/ChatTTS/model/dvae.py

@@ -1,155 +0,0 @@
-import math
-from einops import rearrange
-from vector_quantize_pytorch import GroupedResidualFSQ
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-class ConvNeXtBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        intermediate_dim: int,
-        kernel, dilation,
-        layer_scale_init_value: float = 1e-6,
-    ):
-        # ConvNeXt Block copied from Vocos.
-        super().__init__()
-        self.dwconv = nn.Conv1d(dim, dim, 
-                                kernel_size=kernel, padding=dilation*(kernel//2), 
-                                dilation=dilation, groups=dim
-                            )  # depthwise conv
-        
-        self.norm = nn.LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
-        self.act = nn.GELU()
-        self.pwconv2 = nn.Linear(intermediate_dim, dim)
-        self.gamma = (
-            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
-            if layer_scale_init_value > 0
-            else None
-        )
-
-    def forward(self, x: torch.Tensor, cond = None) -> torch.Tensor:
-        residual = x
-        x = self.dwconv(x)
-        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
-        x = self.norm(x)
-        x = self.pwconv1(x)
-        x = self.act(x)
-        x = self.pwconv2(x)
-        if self.gamma is not None:
-            x = self.gamma * x
-        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
-
-        x = residual + x
-        return x
-    
-
-
-class GFSQ(nn.Module):
-
-    def __init__(self, 
-            dim, levels, G, R, eps=1e-5, transpose = True
-        ):
-        super(GFSQ, self).__init__()
-        self.quantizer = GroupedResidualFSQ(
-            dim=dim,
-            levels=levels,
-            num_quantizers=R,
-            groups=G,
-        )
-        self.n_ind = math.prod(levels)
-        self.eps = eps
-        self.transpose = transpose
-        self.G = G
-        self.R = R
-        
-    def _embed(self, x):
-        if self.transpose:
-            x = x.transpose(1,2)
-        x = rearrange(
-            x, "b t (g r) -> g b t r", g = self.G, r = self.R,
-        )  
-        feat = self.quantizer.get_output_from_indices(x)
-        return feat.transpose(1,2) if self.transpose else feat
-        
-    def forward(self, x,):
-        if self.transpose:
-            x = x.transpose(1,2)
-        feat, ind = self.quantizer(x)
-        ind = rearrange(
-            ind, "g b t r ->b t (g r)",
-        )  
-        embed_onehot = F.one_hot(ind.long(), self.n_ind).to(x.dtype)
-        e_mean = torch.mean(embed_onehot, dim=[0,1])
-        e_mean = e_mean / (e_mean.sum(dim=1) + self.eps).unsqueeze(1)
-        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + self.eps), dim=1))
-        
-        return (
-            torch.zeros(perplexity.shape, dtype=x.dtype, device=x.device),
-            feat.transpose(1,2) if self.transpose else feat,
-            perplexity,
-            None,
-            ind.transpose(1,2) if self.transpose else ind,
-        )
-        
-class DVAEDecoder(nn.Module):
-    def __init__(self, idim, odim,
-                 n_layer = 12, bn_dim = 64, hidden = 256, 
-                 kernel = 7, dilation = 2, up = False
-                ):
-        super().__init__()
-        self.up = up
-        self.conv_in = nn.Sequential(
-            nn.Conv1d(idim, bn_dim, 3, 1, 1), nn.GELU(),
-            nn.Conv1d(bn_dim, hidden, 3, 1, 1)
-        )
-        self.decoder_block = nn.ModuleList([
-            ConvNeXtBlock(hidden, hidden* 4, kernel, dilation,)
-            for _ in range(n_layer)])
-        self.conv_out = nn.Conv1d(hidden, odim, kernel_size=1, bias=False)
-
-    def forward(self, input, conditioning=None):
-        # B, T, C
-        x = input.transpose(1, 2)
-        x = self.conv_in(x)
-        for f in self.decoder_block:
-            x = f(x, conditioning)
-        
-        x = self.conv_out(x)
-        return x.transpose(1, 2)
-    
-
-class DVAE(nn.Module):
-    def __init__(
-        self, decoder_config, vq_config, dim=512
-    ):
-        super().__init__()
-        self.register_buffer('coef', torch.randn(1, 100, 1))
-
-        self.decoder = DVAEDecoder(**decoder_config)
-        self.out_conv = nn.Conv1d(dim, 100, 3, 1, 1, bias=False)
-        if vq_config is not None:
-            self.vq_layer = GFSQ(**vq_config)
-        else:
-            self.vq_layer = None
-
-    def forward(self, inp):
-
-        if self.vq_layer is not None:
-            vq_feats = self.vq_layer._embed(inp)
-        else:
-            vq_feats = inp.detach().clone()
-            
-        temp = torch.chunk(vq_feats, 2, dim=1) # flatten trick :)
-        temp = torch.stack(temp, -1)
-        vq_feats = temp.reshape(*temp.shape[:2], -1)
-
-        vq_feats = vq_feats.transpose(1, 2)
-        dec_out = self.decoder(input=vq_feats)
-        dec_out = self.out_conv(dec_out.transpose(1, 2))
-        mel = dec_out * self.coef
-
-        return mel

xinference/thirdparty/ChatTTS/model/gpt.py

@@ -1,265 +0,0 @@
-import os
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-import logging
-from tqdm import tqdm
-from einops import rearrange
-from transformers.cache_utils import Cache
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.nn.utils.parametrize as P
-from torch.nn.utils.parametrizations import weight_norm
-from transformers import LlamaModel, LlamaConfig
-    
-    
-class LlamaMLP(nn.Module):
-    def __init__(self, hidden_size, intermediate_size):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = F.silu
-
-    def forward(self, x):
-        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-    
-    
-class GPT_warpper(nn.Module):
-    def __init__(
-        self, 
-        gpt_config, 
-        num_audio_tokens,
-        num_text_tokens,
-        num_vq=4,
-        **kwargs,
-        ):
-        super().__init__()
-
-        self.logger = logging.getLogger(__name__)
-        self.gpt = self.build_model(gpt_config)
-        self.model_dim = self.gpt.config.hidden_size 
-
-        self.num_vq = num_vq
-        self.emb_code = nn.ModuleList([nn.Embedding(num_audio_tokens, self.model_dim) for i in range(self.num_vq)])
-        self.emb_text = nn.Embedding(num_text_tokens, self.model_dim)
-        self.head_text = weight_norm(nn.Linear(self.model_dim, num_text_tokens, bias=False), name='weight')
-        self.head_code = nn.ModuleList([weight_norm(nn.Linear(self.model_dim, num_audio_tokens, bias=False), name='weight') for i in range(self.num_vq)])
-
-    def build_model(self, config):
-        
-        configuration = LlamaConfig(**config)
-        model = LlamaModel(configuration)
-        del model.embed_tokens
-        
-        return model
-    
-    def get_emb(self, input_ids, text_mask, **kwargs):
-
-        emb_text = self.emb_text(input_ids[text_mask][:, 0])
-        
-        emb_code = [self.emb_code[i](input_ids[~text_mask][:, i]) for i in range(self.num_vq)]
-        emb_code = torch.stack(emb_code, 2).sum(2)
-        
-        emb = torch.zeros((input_ids.shape[:-1])+(emb_text.shape[-1],), device=emb_text.device, dtype=emb_text.dtype)
-        emb[text_mask] = emb_text
-        emb[~text_mask] = emb_code.to(emb.dtype)
-        
-        return emb
-    
-    def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
-    ):
-        # With static cache, the `past_key_values` is None
-        # TODO joao: standardize interface for the different Cache classes and remove of this if
-        has_static_cache = False
-        if past_key_values is None:
-            past_key_values = getattr(self.gpt.layers[0].self_attn, "past_key_value", None)
-            has_static_cache = past_key_values is not None
-
-        past_length = 0
-        if past_key_values is not None:
-            if isinstance(past_key_values, Cache):
-                past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
-                max_cache_length = (
-                    torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
-                    if past_key_values.get_max_length() is not None
-                    else None
-                )
-                cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
-            # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
-            else:
-                cache_length = past_length = past_key_values[0][0].shape[2]
-                max_cache_length = None
-
-            # Keep only the unprocessed tokens:
-            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
-            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
-            # input)
-            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
-                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
-            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
-            # input_ids based on the past_length.
-            elif past_length < input_ids.shape[1]:
-                input_ids = input_ids[:, past_length:]
-            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-
-            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
-            if (
-                max_cache_length is not None
-                and attention_mask is not None
-                and cache_length + input_ids.shape[1] > max_cache_length
-            ):
-                attention_mask = attention_mask[:, -max_cache_length:]
-
-        position_ids = kwargs.get("position_ids", None)
-        if attention_mask is not None and position_ids is None:
-            # create position_ids on the fly for batch generation
-            position_ids = attention_mask.long().cumsum(-1) - 1
-            position_ids.masked_fill_(attention_mask == 0, 1)
-            if past_key_values:
-                position_ids = position_ids[:, -input_ids.shape[1] :]
-
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and past_key_values is None:
-            model_inputs = {"inputs_embeds": inputs_embeds}
-        else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {"input_ids": input_ids.contiguous()}
-
-        input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
-        if cache_position is None:
-            cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
-        else:
-            cache_position = cache_position[-input_length:]
-
-        if has_static_cache:
-            past_key_values = None
-
-        model_inputs.update(
-            {
-                "position_ids": position_ids,
-                "cache_position": cache_position,
-                "past_key_values": past_key_values,
-                "use_cache": kwargs.get("use_cache"),
-                "attention_mask": attention_mask,
-            }
-        )
-        return model_inputs
-    
-    def generate(
-        self, 
-        emb, 
-        inputs_ids, 
-        temperature, 
-        eos_token, 
-        attention_mask = None,
-        max_new_token = 2048, 
-        min_new_token = 0,
-        LogitsWarpers = [],
-        LogitsProcessors = [],
-        infer_text=False,
-        return_attn=False,
-        return_hidden=False,
-    ):
-        
-        with torch.no_grad():   
-        
-            attentions = []
-            hiddens = []
-            
-            start_idx, end_idx = inputs_ids.shape[1], torch.zeros(inputs_ids.shape[0], device=inputs_ids.device, dtype=torch.long)
-            finish = torch.zeros(inputs_ids.shape[0], device=inputs_ids.device).bool()
-            
-            temperature = temperature[None].expand(inputs_ids.shape[0], -1)
-            temperature = rearrange(temperature, "b n -> (b n) 1")
-
-            attention_mask_cache = torch.ones((inputs_ids.shape[0], inputs_ids.shape[1]+max_new_token,), dtype=torch.bool, device=inputs_ids.device)
-            if attention_mask is not None:
-                attention_mask_cache[:, :attention_mask.shape[1]] = attention_mask
-            
-            for i in tqdm(range(max_new_token)):
-        
-                model_input = self.prepare_inputs_for_generation(inputs_ids, 
-                    outputs.past_key_values if i!=0 else None, 
-                    attention_mask_cache[:, :inputs_ids.shape[1]], use_cache=True)
-            
-                if i == 0:
-                    model_input['inputs_embeds'] = emb
-                else:
-                    if infer_text:
-                        model_input['inputs_embeds'] = self.emb_text(model_input['input_ids'][:,:,0])
-                    else:
-                        code_emb = [self.emb_code[i](model_input['input_ids'][:,:,i]) for i in range(self.num_vq)]
-                        model_input['inputs_embeds'] = torch.stack(code_emb, 3).sum(3)
-                
-                model_input['input_ids'] = None
-                outputs = self.gpt.forward(**model_input, output_attentions=return_attn)
-                attentions.append(outputs.attentions)
-                hidden_states = outputs[0] # 🐻
-                if return_hidden:
-                    hiddens.append(hidden_states[:, -1])
-
-                with P.cached():
-                    if infer_text:
-                        logits = self.head_text(hidden_states) 
-                    else:
-                        logits = torch.stack([self.head_code[i](hidden_states) for i in range(self.num_vq)], 3)
-        
-                logits = logits[:, -1].float()
-
-                if not infer_text:
-                    logits = rearrange(logits, "b c n -> (b n) c")
-                    logits_token = rearrange(inputs_ids[:, start_idx:], "b c n -> (b n) c")
-                else:
-                    logits_token = inputs_ids[:, start_idx:, 0]
-                    
-                logits = logits / temperature
-                
-                for logitsProcessors in LogitsProcessors:
-                    logits = logitsProcessors(logits_token, logits)
-                    
-                for logitsWarpers in LogitsWarpers:
-                    logits = logitsWarpers(logits_token, logits)
-                    
-                if i < min_new_token:
-                    logits[:, eos_token] = -torch.inf
-                
-                scores = F.softmax(logits, dim=-1)
-            
-                idx_next = torch.multinomial(scores, num_samples=1)
-                
-                if not infer_text:
-                    idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq)
-                    finish = finish | (idx_next == eos_token).any(1)
-                    inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(1)], 1)
-                else:
-                    finish = finish | (idx_next == eos_token).any(1)
-                    inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(-1).expand(-1, -1, self.num_vq)], 1)
-
-                end_idx = end_idx + (~finish).int()
-            
-                if finish.all():
-                    break
-            
-            inputs_ids = [inputs_ids[idx, start_idx: start_idx+i] for idx, i in enumerate(end_idx.int())]
-            inputs_ids = [i[:, 0] for i in inputs_ids] if infer_text else inputs_ids
-            
-            if return_hidden:
-                hiddens = torch.stack(hiddens, 1)
-                hiddens = [hiddens[idx, :i] for idx, i in enumerate(end_idx.int())]
-                    
-            if not finish.all():
-                self.logger.warn(f'Incomplete result. hit max_new_token: {max_new_token}')    
-                   
-            return {
-                'ids': inputs_ids, 
-                'attentions': attentions,
-                'hiddens':hiddens,
-            }

xinference/thirdparty/ChatTTS/utils/gpu_utils.py

@@ -1,23 +0,0 @@
-
-import torch
-import logging
-
-def select_device(min_memory = 2048):
-    logger = logging.getLogger(__name__)
-    if torch.cuda.is_available():
-        available_gpus = []
-        for i in range(torch.cuda.device_count()):
-            props = torch.cuda.get_device_properties(i)
-            free_memory = props.total_memory - torch.cuda.memory_reserved(i)
-            available_gpus.append((i, free_memory))
-        selected_gpu, max_free_memory = max(available_gpus, key=lambda x: x[1])
-        device = torch.device(f'cuda:{selected_gpu}')
-        free_memory_mb = max_free_memory / (1024 * 1024)
-        if free_memory_mb < min_memory:
-            logger.log(logging.WARNING, f'GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left.')
-            device = torch.device('cpu')
-    else:
-        logger.log(logging.WARNING, f'No GPU found, use CPU instead')
-        device = torch.device('cpu')
-    
-    return device