Update app.py

app.py

@@ -1,22 +1,36 @@
-import spaces
-from snac import SNAC
+import os
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import snapshot_download
 from dotenv import load_dotenv
+
+# Load environment variables
 load_dotenv()
 
-# Check if CUDA is available
+# Device and torch dtype selection
 device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+# Define a no-op decorator for CPU if needed
+def gpu_decorator(func):
+    return func
+
+# If you are on GPU and have the spaces module, you could replace gpu_decorator with spaces.GPU
+# For CPU usage we simply use a no-op
+# Example: from snac import spaces; gpu_decorator = spaces.GPU()
+
+# Import SNAC after setting device
+from snac import SNAC
 
 print("Loading SNAC model...")
 snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
 snac_model = snac_model.to(device)
+snac_model.eval()  # set SNAC to eval mode
 
 model_name = "canopylabs/orpheus-3b-0.1-ft"
 
-# Download only model config and safetensors
+# Download only model config and safetensors files
 snapshot_download(
     repo_id=model_name,
     allow_patterns=[
@@ -38,36 +52,34 @@ snapshot_download(
     ]
 )
 
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+print("Loading Orpheus model...")
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch_dtype)
 model.to(device)
+model.eval()  # set Orpheus to eval mode
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 print(f"Orpheus model loaded to {device}")
 
-# Process text prompt
+# Process text prompt into tokens with start/end markers
 def process_prompt(prompt, voice, tokenizer, device):
     prompt = f"{voice}: {prompt}"
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    
-    start_token = torch.tensor([[128259]], dtype=torch.int64)  # Start of human
-    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # End of text, End of human
-    
-    modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)  # SOH SOT Text EOT EOH
-    
-    # No padding needed for single input
+
+    start_token = torch.tensor([[128259]], dtype=torch.int64)  # Start token
+    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # End tokens
+
+    modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
     attention_mask = torch.ones_like(modified_input_ids)
-    
     return modified_input_ids.to(device), attention_mask.to(device)
 
-# Parse output tokens to audio
+# Parse output tokens to extract audio codes
 def parse_output(generated_ids):
     token_to_find = 128257
     token_to_remove = 128258
-    
-    token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)
 
+    token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)
     if len(token_indices[1]) > 0:
         last_occurrence_idx = token_indices[1][-1].item()
-        cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
+        cropped_tensor = generated_ids[:, last_occurrence_idx + 1:]
     else:
         cropped_tensor = generated_ids
 
@@ -83,47 +95,43 @@ def parse_output(generated_ids):
         trimmed_row = row[:new_length]
         trimmed_row = [t - 128266 for t in trimmed_row]
         code_lists.append(trimmed_row)
-        
-    return code_lists[0]  # Return just the first one for single sample
 
-# Redistribute codes for audio generation
+    return code_lists[0]  # Return first sample
+
+# Redistribute codes for audio generation using SNAC
 def redistribute_codes(code_list, snac_model):
-    device = next(snac_model.parameters()).device  # Get the device of SNAC model
-    
-    layer_1 = []
-    layer_2 = []
-    layer_3 = []
-    for i in range((len(code_list)+1)//7):
-        layer_1.append(code_list[7*i])
-        layer_2.append(code_list[7*i+1]-4096)
-        layer_3.append(code_list[7*i+2]-(2*4096))
-        layer_3.append(code_list[7*i+3]-(3*4096))
-        layer_2.append(code_list[7*i+4]-(4*4096))
-        layer_3.append(code_list[7*i+5]-(5*4096))
-        layer_3.append(code_list[7*i+6]-(6*4096))
-        
-    # Move tensors to the same device as the SNAC model
+    snac_device = next(snac_model.parameters()).device
+    layer_1, layer_2, layer_3 = [], [], []
+    for i in range((len(code_list) + 1) // 7):
+        layer_1.append(code_list[7 * i])
+        layer_2.append(code_list[7 * i + 1] - 4096)
+        layer_3.append(code_list[7 * i + 2] - (2 * 4096))
+        layer_3.append(code_list[7 * i + 3] - (3 * 4096))
+        layer_2.append(code_list[7 * i + 4] - (4 * 4096))
+        layer_3.append(code_list[7 * i + 5] - (5 * 4096))
+        layer_3.append(code_list[7 * i + 6] - (6 * 4096))
+
     codes = [
-        torch.tensor(layer_1, device=device).unsqueeze(0),
-        torch.tensor(layer_2, device=device).unsqueeze(0),
-        torch.tensor(layer_3, device=device).unsqueeze(0)
+        torch.tensor(layer_1, device=snac_device).unsqueeze(0),
+        torch.tensor(layer_2, device=snac_device).unsqueeze(0),
+        torch.tensor(layer_3, device=snac_device).unsqueeze(0)
     ]
-    
+
     audio_hat = snac_model.decode(codes)
-    return audio_hat.detach().squeeze().cpu().numpy()  # Always return CPU numpy array
+    return audio_hat.detach().squeeze().cpu().numpy()
 
-# Main generation function
-@spaces.GPU()
+# Main generation function with CPU optimizations
+@gpu_decorator
 def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens, progress=gr.Progress()):
     if not text.strip():
         return None
-    
+
     try:
         progress(0.1, "Processing text...")
         input_ids, attention_mask = process_prompt(text, voice, tokenizer, device)
-        
+
         progress(0.3, "Generating speech tokens...")
-        with torch.no_grad():
+        with torch.inference_mode():
             generated_ids = model.generate(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -135,39 +143,38 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
                 num_return_sequences=1,
                 eos_token_id=128258,
             )
-        
+
         progress(0.6, "Processing speech tokens...")
         code_list = parse_output(generated_ids)
-        
-        progress(0.8, "Converting to audio...")
+
+        progress(0.8, "Converting tokens to audio...")
         audio_samples = redistribute_codes(code_list, snac_model)
-        
-        return (24000, audio_samples)  # Return sample rate and audio
+
+        return (24000, audio_samples)  # Return sample rate and numpy array audio
     except Exception as e:
         print(f"Error generating speech: {e}")
         return None
 
-# Examples for the UI
+# Example inputs for the Gradio UI
 examples = [
     ["Hey there my name is Tara, <chuckle> and I'm a speech generation model that can sound like a person.", "tara", 0.6, 0.95, 1.1, 1200],
     ["I've also been taught to understand and produce paralinguistic things like sighing, or chuckling, or yawning!", "dan", 0.7, 0.95, 1.1, 1200],
-    ["I live in San Francisco, and have, uhm let's see, 3 billion 7 hundred ... well, lets just say a lot of parameters.", "emma", 0.6, 0.9, 1.2, 1200]
+    ["I live in San Francisco, and have, uhm let's see, 3 billion 7 hundred ... well, let's just say a lot of parameters.", "emma", 0.6, 0.9, 1.2, 1200]
 ]
 
-# Available voices
 VOICES = ["tara", "dan", "josh", "emma"]
 
 # Create Gradio interface
 with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
     gr.Markdown("""
-    # 🎵 [Orpheus Text-to-Speech](https://github.com/canopyai/Orpheus-TTS)
-    Enter your text below and hear it converted to natural-sounding speech with the Orpheus TTS model.
+    # 🎵 Orpheus Text-to-Speech
+    Enter your text below and hear it converted to natural-sounding speech.
     
-    ## Tips for better prompts:
-    - Add paralinguistic elements like `<chuckle>`, `<sigh>`, or `uhm` for more human-like speech.
-    - Longer text prompts generally work better than very short phrases
-    - Adjust the temperature slider for more varied (higher) or consistent (lower) speech patterns
-    """)    
+    **Tips for better prompts:**
+    - Include paralinguistic elements like `<chuckle>`, `<sigh>`, or `uhm` for more human-like speech.
+    - Longer prompts often produce more natural results.
+    - Adjust the temperature slider to control variation in speech patterns.
+    """)
     with gr.Row():
         with gr.Column(scale=3):
             text_input = gr.Textbox(
@@ -180,37 +187,33 @@ with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
                 value="tara", 
                 label="Voice"
             )
-            
             with gr.Accordion("Advanced Settings", open=False):
                 temperature = gr.Slider(
                     minimum=0.1, maximum=1.5, value=0.6, step=0.05,
-                    label="Temperature", 
+                    label="Temperature",
                     info="Higher values (0.7-1.0) create more expressive but less stable speech"
                 )
                 top_p = gr.Slider(
                     minimum=0.1, maximum=1.0, value=0.95, step=0.05,
-                    label="Top P", 
+                    label="Top P",
                     info="Nucleus sampling threshold"
                 )
                 repetition_penalty = gr.Slider(
                     minimum=1.0, maximum=2.0, value=1.1, step=0.05,
-                    label="Repetition Penalty", 
+                    label="Repetition Penalty",
                     info="Higher values discourage repetitive patterns"
                 )
                 max_new_tokens = gr.Slider(
                     minimum=100, maximum=2000, value=1200, step=100,
-                    label="Max Length", 
+                    label="Max Length",
                     info="Maximum length of generated audio (in tokens)"
                 )
-            
             with gr.Row():
                 submit_btn = gr.Button("Generate Speech", variant="primary")
                 clear_btn = gr.Button("Clear")
-                
         with gr.Column(scale=2):
             audio_output = gr.Audio(label="Generated Speech", type="numpy")
-            
-    # Set up examples
+    
     gr.Examples(
         examples=examples,
         inputs=[text_input, voice, temperature, top_p, repetition_penalty, max_new_tokens],
@@ -219,19 +222,17 @@ with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
         cache_examples=True,
     )
     
-    # Set up event handlers
     submit_btn.click(
         fn=generate_speech,
         inputs=[text_input, voice, temperature, top_p, repetition_penalty, max_new_tokens],
         outputs=audio_output
     )
-    
     clear_btn.click(
         fn=lambda: (None, None),
         inputs=[],
         outputs=[text_input, audio_output]
     )
 
-# Launch the app
+# Launch the Gradio app
 if __name__ == "__main__":
-    demo.queue().launch(share=False, ssr_mode=False)
+    demo.queue().launch(share=False, ssr_mode=False)