Update app.py

app.py

@@ -1,102 +1,157 @@
-from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
-from PIL import Image
-import requests
-import torch
-from threading import Thread
 import gradio as gr
-from gradio import FileData
-import time
 import spaces
+import torch
+from transformers import AutoTokenizer,VitsModel
 import os
+import numpy as np  
+
 token=os.environ.get("key_")
-ckpt = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-model = MllamaForConditionalGeneration.from_pretrained(ckpt,token=token,
-    torch_dtype=torch.bfloat16).to("cuda")
-processor = AutoProcessor.from_pretrained(ckpt,token=token)
+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 bot_streaming(message, history, max_new_tokens=250):
-    
-    txt = message["text"]
-    ext_buffer = f"{txt}"
-    
-    messages= [] 
-    images = []
-    
+def  get_model(name_model):
+    global models
+    if name_model in   models:
+        return models[name_model]
+    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)
+    return models[name_model]
 
-    for i, msg in enumerate(history): 
-        if isinstance(msg[0], tuple):
-            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
-            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
-            images.append(Image.open(msg[0][0]).convert("RGB"))
-        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
-            # messages are already handled
-            pass
-        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
-            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
-            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
-
-    # add current message
-    if len(message["files"]) == 1:
-        
-        if isinstance(message["files"][0], str): # examples
-            image = Image.open(message["files"][0]).convert("RGB")
-        else: # regular input
-            image = Image.open(message["files"][0]["path"]).convert("RGB")
-        images.append(image)
-        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
-    else:
-        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
-
-
-    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
-
-    if images == []:
-        inputs = processor(text=texts, return_tensors="pt").to("cuda")
-    else:
-        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
-    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
-
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
-    generated_text = ""
-    
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    buffer = ""
-    
-    for new_text in streamer:
-        buffer += new_text
-        generated_text_without_prompt = buffer
-        time.sleep(0.01)
-        yield buffer
-
-
-demo = gr.ChatInterface(fn=bot_streaming, title="Multimodal Llama", examples=[
-    [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]},
-    200],
-    [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]},
-    250],
-    [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]},
-    250],
-    [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]},
-    250],
-    [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 
-    250],
-    ],
-      textbox=gr.MultimodalTextbox(), 
-      additional_inputs = [gr.Slider(
-              minimum=10,
-              maximum=500,
-              value=250,
-              step=10,
-              label="Maximum number of new tokens to generate",
-          )
-        ],
-      cache_examples=False,
-      description="Try Multimodal Llama by Meta with transformers in this demo. Upload an image, and start chatting about it, or simply try one of the examples below. To learn more about Llama Vision, visit [our blog post](https://huggingface.co/blog/llama32). ",
-      stop_btn="Stop Generation", 
-      fill_height=True,
-    multimodal=True)
+
+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):
+     
     
-demo.launch(debug=True)
+     inputs = tokenizer(text, return_tensors="pt")
+     model=get_model(name_model)
+     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"
+                                
+                                  
+                            ],
+                            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()