app.py

import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer,VitsModel
import os
import numpy as np  

token=os.environ.get("key_")
#tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vtk",token=token)
models= {}

import noisereduce as nr

import torch
from typing import Any, Callable, Optional, Tuple, Union,Iterator

import torch.nn as nn # Import the missing module
def remove_noise_nr(audio_data,sr=16000):
    reduced_noise = nr.reduce_noise(y=audio_data,hop_length=256, sr=sr)
    return reduced_noise

def _inference_forward_stream(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        speaker_embeddings: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        padding_mask: Optional[torch.Tensor] = None,
        chunk_size: int = 32,  # Chunk size for streaming output
        is_streaming: bool = True,
    ) -> Iterator[torch.Tensor]:
        """Generates speech waveforms in a streaming fashion."""
        if attention_mask is not None:
            padding_mask = attention_mask.unsqueeze(-1).float()
        else:
            padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()



        text_encoder_output = self.text_encoder(
            input_ids=input_ids,
            padding_mask=padding_mask,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
        hidden_states = hidden_states.transpose(1, 2)
        input_padding_mask = padding_mask.transpose(1, 2)

        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances

        if self.config.use_stochastic_duration_prediction:
            log_duration = self.duration_predictor(
                hidden_states,
                input_padding_mask,
                speaker_embeddings,
                reverse=True,
                noise_scale=self.noise_scale_duration,
            )
        else:
            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)

        length_scale = 1.0 / self.speaking_rate
        duration = torch.ceil(torch.exp(log_duration) * input_padding_mask * length_scale)
        predicted_lengths = torch.clamp_min(torch.sum(duration, [1, 2]), 1).long()


        # Create a padding mask for the output lengths of shape (batch, 1, max_output_length)
        indices = torch.arange(predicted_lengths.max(), dtype=predicted_lengths.dtype, device=predicted_lengths.device)
        output_padding_mask = indices.unsqueeze(0) < predicted_lengths.unsqueeze(1)
        output_padding_mask = output_padding_mask.unsqueeze(1).to(input_padding_mask.dtype)

        # Reconstruct an attention tensor of shape (batch, 1, out_length, in_length)
        attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(output_padding_mask, -1)
        batch_size, _, output_length, input_length = attn_mask.shape
        cum_duration = torch.cumsum(duration, -1).view(batch_size * input_length, 1)
        indices = torch.arange(output_length, dtype=duration.dtype, device=duration.device)
        valid_indices = indices.unsqueeze(0) < cum_duration
        valid_indices = valid_indices.to(attn_mask.dtype).view(batch_size, input_length, output_length)
        padded_indices = valid_indices - nn.functional.pad(valid_indices, [0, 0, 1, 0, 0, 0])[:, :-1]
        attn = padded_indices.unsqueeze(1).transpose(2, 3) * attn_mask

        # Expand prior distribution
        prior_means = torch.matmul(attn.squeeze(1), prior_means).transpose(1, 2)
        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances).transpose(1, 2)

        prior_latents = prior_means + torch.randn_like(prior_means) * torch.exp(prior_log_variances) * self.noise_scale
        latents = self.flow(prior_latents, output_padding_mask, speaker_embeddings, reverse=True)

        spectrogram = latents * output_padding_mask
        if is_streaming:
            
            for i in range(0, spectrogram.size(-1), chunk_size):
                with torch.no_grad():
                    wav=self.decoder(spectrogram[:,:,i : i + chunk_size] ,speaker_embeddings)
                yield wav.squeeze().cpu().numpy()
        else:

              wav=self.decoder(spectrogram,speaker_embeddings)
              yield  wav.squeeze().cpu().numpy()
@spaces.GPU
def  get_model(name_model):
    global models
    if name_model in   models:
        if  name_model=='wasmdashai/vits-en-v1':
            tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vits-en-v1",token=token)
        else: 
            tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vtk",token=token)

        

        
        return models[name_model],tokenizer
    models[name_model]=VitsModel.from_pretrained(name_model,token=token).cuda()
    
    
   
    models[name_model].decoder.apply_weight_norm()
    # torch.nn.utils.weight_norm(self.decoder.conv_pre)
    # torch.nn.utils.weight_norm(self.decoder.conv_post)
    for flow in models[name_model].flow.flows:
        torch.nn.utils.weight_norm(flow.conv_pre)
        torch.nn.utils.weight_norm(flow.conv_post)
    
    if  name_model=='wasmdashai/vits-en-v1':
            tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vits-en-v1",token=token)
    else: 
            tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vtk",token=token)

    return models[name_model],tokenizer


zero = torch.Tensor([0]).cuda()
print(zero.device) # <-- 'cpu' 🤔
import torch
TXT="""السلام  عليكم  ورحمة الله وبركاتة  يا هلا وسهلا ومراحب بالغالي  اخباركم  طيبين ان شاء الله     ارحبوا  على العين والراس     """
@spaces.GPU
def   modelspeech(text=TXT,name_model="wasmdashai/vits-ar-sa-huba-v2",speaking_rate=16000):
     model,tokenizer=get_model(name_model)
    
     inputs = tokenizer(text, return_tensors="pt")
     
     model.speaking_rate=speaking_rate
     with torch.no_grad():
        wav=list(_inference_forward_stream(model,input_ids=inputs.input_ids.cuda(),attention_mask=inputs.attention_mask.cuda(),speaker_embeddings= None,is_streaming=False))[0]
     # with torch.no_grad():
     #      wav = model(input_ids=inputs["input_ids"].cuda()).waveform.cpu().numpy().reshape(-1)#.detach()
          
     return  (model.config.sampling_rate,wav)

model_choices = gr.Dropdown(
                            choices=[
                                
                                "wasmdashai/vits-ar-sa-huba-v1",
                                 "wasmdashai/vits-ar-sa-huba-v2",
                                 
                                 "wasmdashai/vits-ar-sa-A",
                                "wasmdashai/vits-ar-ye-sa",
                                "wasmdashai/vits-ar-sa-M-v1",
                                'wasmdashai/vits-en-v1'
                                
                                  
                            ],
                            label="اختر النموذج",
                            value="wasmdashai/vits-ar-sa-huba-v2",
                        )

demo = gr.Interface(fn=modelspeech, inputs=["text",model_choices,gr.Slider(0.1, 1, step=0.1,value=0.8)], outputs=["audio"])
demo.queue()
demo.launch()