Update Bark.py

开源语音模型 Bark/Bark.py

@@ -9,7 +9,7 @@ st.set_page_config(
     layout='wide'
 )
 
-choose_load_method = 0
+choose_load_method = 1
 
 if choose_load_method == 0:
     # GPU加载（需要5G显存）
@@ -22,7 +22,17 @@ if choose_load_method == 0:
     model, processor = load_bark_model()
 
 elif choose_load_method == 1:
-    # CPU加载（需要9G内存，运行速度慢，不推荐）
+    # GPU加载bark-small模型（需要3G显存）
+    @st.cache_resource
+    def load_bark_model():
+        from transformers import AutoProcessor, AutoModel
+        processor = AutoProcessor.from_pretrained("suno/bark-small")
+        model = AutoModel.from_pretrained("suno/bark-small").to("cuda")
+        return model, processor
+    model, processor = load_bark_model()
+
+elif choose_load_method == 2:
+    # CPU加载bark模型（需要9G内存，运行速度慢，不推荐）
     @st.cache_resource
     def load_bark_model():
         from transformers import AutoProcessor, AutoModel
@@ -30,6 +40,16 @@ elif choose_load_method == 1:
         model = AutoModel.from_pretrained("suno/bark")
         return model, processor
     model, processor = load_bark_model()
+
+elif choose_load_method == 3:
+    # CPU加载bark-small模型（需要5G内存，运行速度慢，不推荐）
+    @st.cache_resource
+    def load_bark_model():
+        from transformers import AutoProcessor, AutoModel
+        processor = AutoProcessor.from_pretrained("suno/bark-small")
+        model = AutoModel.from_pretrained("suno/bark-small")
+        return model, processor
+    model, processor = load_bark_model()
     
 prompt = st.chat_input("在这里输入您的命令")
 
@@ -48,7 +68,7 @@ if prompt:
         text=[prompt],
         return_tensors="pt",
     )
-    if choose_load_method == 0:
+    if choose_load_method == 0 or choose_load_method == 1:
         inputs = {key: value.to("cuda") for key, value in inputs.items()}
 
     speech_values = model.generate(**inputs, do_sample=True)