Update app.py

app.py

@@ -13,10 +13,13 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
 repo_id =  "parler-tts/parler-tts-mini-v1"
-repo_id_large = "ylacombe/parler-large-v1-og"
+repo_id_large = "parler-tts/parler-tts-large-v1"
+repo_id_tiny =  "parler-tts/parler-tts-tiny-v1"
 
 model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id).to(device)
 model_large = ParlerTTSForConditionalGeneration.from_pretrained(repo_id_large).to(device)
+model_tiny = ParlerTTSForConditionalGeneration.from_pretrained(repo_id_tiny).to(device)
+
 tokenizer = AutoTokenizer.from_pretrained(repo_id)
 feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
 
@@ -76,19 +79,23 @@ def preprocess(text):
     return text
 
 @spaces.GPU
-def gen_tts(text, description, use_large=False):
+def gen_tts(text, description, version_to_use=False):
     inputs = tokenizer(description.strip(), return_tensors="pt").to(device)
     prompt = tokenizer(preprocess(text), return_tensors="pt").to(device)
 
     set_seed(SEED)
-    if use_large:
+    if version_to_use=="Large":
         generation = model_large.generate(
             input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids, attention_mask=inputs.attention_mask, prompt_attention_mask=prompt.attention_mask, do_sample=True, temperature=1.0
         )      
-    else:
+    elif version_to_use=="Miny":
         generation = model.generate(
             input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids, attention_mask=inputs.attention_mask, prompt_attention_mask=prompt.attention_mask, do_sample=True, temperature=1.0
         )
+    else:
+        generation = model_tiny.generate(
+            input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids, attention_mask=inputs.attention_mask, prompt_attention_mask=prompt.attention_mask, do_sample=True, temperature=1.0
+        )
     audio_arr = generation.cpu().numpy().squeeze()
 
     return SAMPLE_RATE, audio_arr
@@ -163,12 +170,12 @@ with gr.Blocks(css=css) as block:
         with gr.Column():
             input_text = gr.Textbox(label="Input Text", lines=2, value=default_text, elem_id="input_text")
             description = gr.Textbox(label="Description", lines=2, value=default_description, elem_id="input_description")
-            use_large = gr.Checkbox(value=False, label="Use Large checkpoint", info="Generate with Parler-TTS Large v1 instead of Mini v1 - Better but way slower.")
+            version_to_use = gr.Radio(["Tiny", "Mini", "Large"], value="Mini", label="Checkpoint to use", info="The larger the model, the better it is, at the cost of speed.")
             run_button = gr.Button("Generate Audio", variant="primary")
         with gr.Column():
             audio_out = gr.Audio(label="Parler-TTS generation", type="numpy", elem_id="audio_out")
 
-    inputs = [input_text, description, use_large]
+    inputs = [input_text, description, version_to_use]
     outputs = [audio_out]
     run_button.click(fn=gen_tts, inputs=inputs, outputs=outputs, queue=True)
     gr.Examples(examples=examples, fn=gen_tts, inputs=inputs, outputs=outputs, cache_examples=True)