Update app.py

app.py

@@ -7,21 +7,46 @@ import os
 
 token=os.environ.get("key_")
 tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vtk",token=token)
+models= {}
+
+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]
+
 
-model=VitsModel.from_pretrained("wasmdashai/vtk",token=token).cuda()
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
 import torch
 @spaces.GPU
-def   modelspeech(text):
+def   modelspeech(text,name_model):
      
     
      inputs = tokenizer(text, return_tensors="pt")
+     model=get_model(name_model)
      with torch.no_grad():
           wav = model(input_ids=inputs["input_ids"].cuda()).waveform.cpu().numpy().reshape(-1)#.detach()
           
      return  model.config.sampling_rate,wav#remove_noise_nr(wav)
 
-
-demo = gr.Interface(fn=modelspeech, inputs=["text"], outputs=["audio"])
+model_choices = gr.Dropdown(
+                            choices=[
+                                "wasmdashai/vits-ar-sa",
+                                "wasmdashai/vits-ar-sa-huba",
+                                "wasmdashai/vits-ar-sa-ms",
+                                "wasmdashai/vits-ar-sa-magd",
+                                "wasmdashai/vtk",
+                            ],
+                            label="اختر النموذج",
+                            value="wasmdashai/vtk",
+                        )
+demo = gr.Interface(fn=modelspeech, inputs=["text",model_choices], outputs=["audio"])
 demo.launch()