Update app.py

app.py

@@ -42,36 +42,51 @@ _model_cache: Dict[str, object] = {}
 _device_hint = "auto"
 
 def _load_bark():
+    import torch
     from transformers import pipeline
-    pipe = pipeline("text-to-speech", model="suno/bark-small", device_map=_device_hint)
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    pipe = pipeline(
+        task="text-to-speech",
+        model="suno/bark-small",
+        device=device,                        
+        model_kwargs={"low_cpu_mem_usage": False, "torch_dtype": torch.float32}
+    )
+
     if getattr(pipe.model.config, "pad_token_id", None) is None:
         pipe.model.config.pad_token_id = pipe.model.config.eos_token_id
 
-    def generate(text: str) -> Tuple[int, np.ndarray]:
+    def generate(text: str):
         out = pipe(text)
         return int(out["sampling_rate"]), np.asarray(out["audio"], dtype=np.float32)
 
     return generate
 
+
 def _load_mms():
+    import torch
     from transformers import pipeline
-    pipe = pipeline("text-to-speech", model="facebook/mms-tts-rus", device_map=_device_hint)
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    pipe = pipeline(
+        "text-to-speech",
+        model="facebook/mms-tts-rus",
+        device=device,
+        model_kwargs={"low_cpu_mem_usage": False, "torch_dtype": torch.float32}
+    )
     if getattr(pipe.model.config, "pad_token_id", None) is None:
         pipe.model.config.pad_token_id = pipe.model.config.eos_token_id
 
-    def generate(text: str) -> Tuple[int, np.ndarray]:
+    def generate(text: str):
         out = pipe(text)
         return int(out["sampling_rate"]), np.asarray(out["audio"], dtype=np.float32)
-
     return generate
 
+
 def _load_seamless():
     import torch
     import numpy as np
     from transformers import AutoProcessor
-    from transformers.models.seamless_m4t_v2.modeling_seamless_m4t_v2 import (
-        SeamlessM4Tv2Model,
-    )
+    from transformers.models.seamless_m4t_v2.modeling_seamless_m4t_v2 import SeamlessM4Tv2Model
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -79,7 +94,11 @@ def _load_seamless():
         "facebook/seamless-m4t-v2-large",
         use_fast=False
     )
-    model = SeamlessM4Tv2Model.from_pretrained("facebook/seamless-m4t-v2-large").to(device)
+    
+    model = SeamlessM4Tv2Model.from_pretrained(
+        "facebook/seamless-m4t-v2-large",
+        low_cpu_mem_usage=False
+    ).to(device)
 
     def generate(text: str):
         inputs = proc(text=text, src_lang="rus", return_tensors="pt")
@@ -88,9 +107,9 @@ def _load_seamless():
             audio = model.generate(**inputs, tgt_lang="rus")[0]
         audio = audio.detach().cpu().numpy().squeeze().astype(np.float32)
         return 16000, audio
-
     return generate
 
+
 def get_generator(kind: str):
     if kind in _model_cache:
         return _model_cache[kind]
@@ -171,13 +190,18 @@ MUSIC_MODELS = [
 ]
 
 def get_music_pipe(model_name: str):
-    if model_name in _music_pipes:
-        return _music_pipes[model_name]
+    import torch
     from transformers import pipeline
-    pipe = pipeline("text-to-audio", model=model_name, device_map=_device_hint)
-    _music_pipes[model_name] = pipe
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    pipe = pipeline(
+        "text-to-audio",
+        model=model_name,
+        device=device,
+        model_kwargs={"low_cpu_mem_usage": False, "torch_dtype": torch.float32}
+    )
     return pipe
 
+
 MUSIC_DEFAULT_PROMPTS = (
     "High-energy 90s rock track with distorted electric guitars, driving bass, and hard-hitting acoustic drums\n"
     "Modern electronic dance track with punchy kick, bright synth lead, and sidechained pads, 128 BPM\n"