Update app.py

app.py

@@ -7,7 +7,6 @@ from huggingface_hub import snapshot_download
 from dotenv import load_dotenv
 load_dotenv()
 
-
 # Check if CUDA is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -59,23 +58,6 @@ def process_prompt(prompt, voice, tokenizer, device):
     
     return modified_input_ids.to(device), attention_mask.to(device)
 
-# Generate speech tokens
-@spaces.GPU()
-def generate_speech_tokens(input_ids, attention_mask, model, params):
-    with torch.no_grad():
-        generated_ids = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            max_new_tokens=params["max_new_tokens"],
-            do_sample=True,
-            temperature=params["temperature"],
-            top_p=params["top_p"],
-            repetition_penalty=params["repetition_penalty"],
-            num_return_sequences=1,
-            eos_token_id=128258,
-        )
-    return generated_ids
-
 # Parse output tokens to audio
 def parse_output(generated_ids):
     token_to_find = 128257
@@ -131,6 +113,7 @@ def redistribute_codes(code_list, snac_model):
     return audio_hat.detach().squeeze().cpu().numpy()  # Always return CPU numpy array
 
 # Main generation function
+@spaces.GPU()
 def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens, progress=gr.Progress()):
     if not text.strip():
         return None
@@ -140,13 +123,18 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
         input_ids, attention_mask = process_prompt(text, voice, tokenizer, device)
         
         progress(0.3, "Generating speech tokens...")
-        params = {
-            "temperature": temperature,
-            "top_p": top_p,
-            "repetition_penalty": repetition_penalty,
-            "max_new_tokens": max_new_tokens
-        }
-        generated_ids = generate_speech_tokens(input_ids, attention_mask, model, params)
+        with torch.no_grad():
+            generated_ids = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                num_return_sequences=1,
+                eos_token_id=128258,
+            )
         
         progress(0.6, "Processing speech tokens...")
         code_list = parse_output(generated_ids)