Update app.py

app.py

@@ -1,43 +1,42 @@
 import gradio as gr
-from parler_tts import ParlerTTSForConditionalGeneration
+from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSConfig
 from transformers import AutoTokenizer, set_seed
 import soundfile as sf
 import torch
 import os
+from accelerate import Accelerator
+from accelerate.utils import set_seed, AutocastKwargs
 
 os.system("bash install.sh")
 
-device = "cpu"
-if torch.cuda.is_available():
-    device = "cuda:0"
-if torch.backends.mps.is_available():
-    device = "mps"
-if torch.xpu.is_available():
-    device = "xpu"
+# Setup device and accelerator
+accelerator = Accelerator()
+device = accelerator.device
+mixed_precision = "no" if device == "cpu" else "bf16"
+torch_dtype = torch.float32 if device == "cpu" else torch.bfloat16
 
-torch_dtype = torch.bfloat16 if device != "cpu" else torch.float32
-model = ParlerTTSForConditionalGeneration.from_pretrained("AkhilTolani/parler-tts-finetune-vocals-only-large-18720-steps", torch_dtype=torch_dtype, attn_implementation="sdpa").to(device)
+# Load model and tokenizer
+model_path = "AkhilTolani/parler-tts-finetune-vocals-only-large-18720-steps"
+config = ParlerTTSConfig.from_pretrained(model_path)
+model = ParlerTTSForConditionalGeneration.from_pretrained(
+    model_path,
+    config=config,
+    torch_dtype=torch_dtype,
+    attn_implementation="sdpa"
+)
+model = accelerator.prepare(model)
 
 tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-xxl")
 
 def generate_audio(prompt, description, seed, temperature, max_length, do_sample):
-    # Set the seed for reproducibility
     seed = int(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-    if torch.backends.mps.is_available():
-        torch.backends.mps.manual_seed(seed)
-    if torch.xpu.is_available():
-        torch.xpu.manual_seed(seed)
+    set_seed(seed)
 
     input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
     prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
-    # Define num_codebooks from the model configuration
     num_codebooks = model.decoder.config.num_codebooks
 
-    # Set up generation arguments
     gen_kwargs = {
         "do_sample": do_sample,
         "temperature": temperature,
@@ -45,22 +44,53 @@ def generate_audio(prompt, description, seed, temperature, max_length, do_sample
         "min_new_tokens": num_codebooks + 1,
     }
 
-    set_seed(seed)
+    # Prepare for generation
+    batch = {
+        "input_ids": input_ids,
+        "prompt_input_ids": prompt_input_ids,
+    }
+
+    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "no"))
+
+    # Generation step
+    def generate_step(batch, accelerator):
+        batch.pop("decoder_attention_mask", None)
+        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=True)
+        
+        with accelerator.autocast(**autocast_kwargs):
+            if mixed_precision == "fp16":
+                encoder_outputs = eval_model.text_encoder(
+                    input_ids=batch.get("input_ids"),
+                    attention_mask=batch.get("attention_mask", None)
+                )
+                encoder_hidden_states = encoder_outputs.last_hidden_state
+                if (
+                    config.text_encoder.hidden_size != config.decoder.hidden_size
+                    and config.decoder.cross_attention_hidden_size is None
+                ):
+                    encoder_hidden_states = eval_model.enc_to_dec_proj(encoder_hidden_states)
+                
+                if batch.get("attention_mask", None) is not None:
+                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]
+                
+                encoder_outputs.last_hidden_state = encoder_hidden_states
+                batch["encoder_outputs"] = encoder_outputs
+
+            output_audios = eval_model.generate(**batch, **gen_kwargs)
+        return output_audios
 
-    # Generate the output
-    generation = model.generate(
-        input_ids=input_ids,
-        prompt_input_ids=prompt_input_ids,
-        **gen_kwargs
-    ).to(torch.float32)
+    with torch.no_grad():
+        generated_audios = generate_step(batch, accelerator)
 
-    audio_arr = generation.cpu().numpy().squeeze()
+    # Post-process the generated audio
+    audio_arr = generated_audios.cpu().numpy().squeeze()
     sf.write("parler_tts_out.wav", audio_arr, model.config.sampling_rate)
 
     return "parler_tts_out.wav"
 
-default_prompt = "free monster when you fucking hoes walk in the room chilettoes sitting up straight my fucking pasha been exposed somoomits start to hurt ive been shitting real gold doing dumb shit you never know or never seen they call me leopard i stay with the green"
-default_description = "Experience the vibrant energy of a hip hop track featuring a male rapper delivering smooth verses over a catchy synth lead melody, supported by punchy kicks, deep 808 bass, claps, and shimmering hi-hats. The song transitions seamlessly to a female vocalist singing melodically alongside a raw synth bass line. Perfect for setting the mood in a solo dance session at home or keeping the party going in a lively club environment."
+# Gradio interface setup
+default_prompt = "thought no beef im hate to get murder right in these streets i told yall niggins is dead fucking green tbs and tsg my shit only you cant beat out if you aint going to aim and squeeze take your mvp out the game just like a referee im talking about my life you just rapping on beats i be clapping on streets theyre using technology to try to find where the bullets coming from they wont find those z nope because im a smooth criminal i got some screwed loose because im a sick of the"
+default_description = "A male vocalist delivers an energetic and passionate freestyle in a medium-fast tempo, showcasing an enthusiastic and emotional performance with emphatic expression, conveying a youthful and groovy vibe throughout the track."
 default_seed = "456"
 
 interface = gr.Interface(
@@ -69,14 +99,14 @@ interface = gr.Interface(
         gr.Textbox(label="Prompt", value=default_prompt), 
         gr.Textbox(label="Description", value=default_description),
         gr.Textbox(label="Seed", value=default_seed),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.1, value=0.75),
+        gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.1, value=1.0),
         gr.Slider(label="Max Length", minimum=256, maximum=5120, step=256, value=2580),
         gr.Dropdown(label="Do Sample", choices=[True, False], value=True)
     ],
     outputs=gr.Audio(label="Generated Audio"),
     title="Parler TTS Audio Generation",
-    description="Generate audio using the Parler TTS model. Provide a prompt, description, and seed to generate the corresponding audio."
+    description="Generate vocals using the Parler TTS model. Provide a prompt, description, and seed to generate the corresponding audio."
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()