Update app.py

app.py

@@ -8,29 +8,28 @@ from dotenv import load_dotenv
 # Load environment variables
 load_dotenv()
 
+# Set number of threads (adjust based on your CPU cores)
+torch.set_num_threads(4)
+
 # 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
+# No-op decorator for CPU mode (if you had GPU-specific decorators)
 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
+snac_model.eval()  # Set SNAC to eval mode
 
 model_name = "canopylabs/orpheus-3b-0.1-ft"
 
-# Download only model config and safetensors files
+# Download only necessary files for the Orpheus model
 snapshot_download(
     repo_id=model_name,
     allow_patterns=[
@@ -55,23 +54,30 @@ snapshot_download(
 print("Loading Orpheus model...")
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch_dtype)
 model.to(device)
-model.eval()  # set Orpheus to eval mode
+model.eval()  # Set the model to evaluation mode
+
+# Optionally compile the model for PyTorch 2.0+ on CPU (if available)
+if hasattr(torch, "compile") and device == "cpu":
+    try:
+        model = torch.compile(model)
+        print("Model compiled with torch.compile")
+    except Exception as e:
+        print("torch.compile not supported:", e)
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 print(f"Orpheus model loaded to {device}")
 
-# 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 token
-    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # End tokens
+    start_token = torch.tensor([[128259]], dtype=torch.int64)
+    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)
 
     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 extract audio codes
 def parse_output(generated_ids):
     token_to_find = 128257
     token_to_remove = 128258
@@ -96,9 +102,8 @@ def parse_output(generated_ids):
         trimmed_row = [t - 128266 for t in trimmed_row]
         code_lists.append(trimmed_row)
 
-    return code_lists[0]  # Return first sample
+    return code_lists[0]
 
-# Redistribute codes for audio generation using SNAC
 def redistribute_codes(code_list, snac_model):
     snac_device = next(snac_model.parameters()).device
     layer_1, layer_2, layer_3 = [], [], []
@@ -116,21 +121,17 @@ def redistribute_codes(code_list, snac_model):
         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()
 
-# 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...")
+        progress(0.05, "Processing text...")
         input_ids, attention_mask = process_prompt(text, voice, tokenizer, device)
-
-        progress(0.3, "Generating speech tokens...")
+        progress(0.2, "Generating tokens...")
         with torch.inference_mode():
             generated_ids = model.generate(
                 input_ids=input_ids,
@@ -143,71 +144,73 @@ 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...")
+        progress(0.4, "Parsing tokens...")
         code_list = parse_output(generated_ids)
-
-        progress(0.8, "Converting tokens to audio...")
+        progress(0.7, "Generating audio...")
         audio_samples = redistribute_codes(code_list, snac_model)
-
-        return (24000, audio_samples)  # Return sample rate and numpy array audio
+        progress(1.0, "Done")
+        return (24000, audio_samples)
     except Exception as e:
         print(f"Error generating speech: {e}")
         return None
 
-# Example inputs for the Gradio UI
+def convert_model_to_onnx():
+    """
+    Converts the Orpheus model to ONNX format using a dummy prompt.
+    The exported file will be saved as 'orpheus_model.onnx' in the working directory.
+    """
+    dummy_prompt = "tara: Hello"
+    dummy_input = tokenizer(dummy_prompt, return_tensors="pt").input_ids.to(device)
+    file_path = "orpheus_model.onnx"
+    try:
+        # Export the model to ONNX format
+        torch.onnx.export(
+            model,
+            dummy_input,
+            file_path,
+            export_params=True,
+            opset_version=14,
+            input_names=["input_ids"],
+            output_names=["logits"],
+            dynamic_axes={
+                "input_ids": {0: "batch_size", 1: "sequence_length"},
+                "logits": {0: "batch_size", 1: "sequence_length"}
+            },
+        )
+        return f"Model converted to ONNX and saved as '{file_path}'."
+    except Exception as e:
+        return f"Error during ONNX conversion: {e}"
+
+# UI examples and voice choices
 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, let's just say a lot of parameters.", "emma", 0.6, 0.9, 1.2, 1200]
 ]
-
 VOICES = ["tara", "dan", "josh", "emma"]
 
-# Create Gradio interface
 with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
     gr.Markdown("""
     # 🎵 Orpheus Text-to-Speech
-    Enter your text below and hear it converted to natural-sounding speech.
+    Enter text to hear it converted to natural-sounding speech.
     
-    **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.
+    **Tips:**
+    - Use paralinguistic cues like `<chuckle>` or `<sigh>`.
+    - Longer text can produce more natural results.
     """)
     with gr.Row():
         with gr.Column(scale=3):
-            text_input = gr.Textbox(
-                label="Text to speak", 
-                placeholder="Enter your text here...",
-                lines=5
-            )
-            voice = gr.Dropdown(
-                choices=VOICES, 
-                value="tara", 
-                label="Voice"
-            )
+            text_input = gr.Textbox(label="Text to speak", placeholder="Enter your text...", lines=5)
+            voice = gr.Dropdown(choices=VOICES, 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",
-                    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",
-                    info="Nucleus sampling threshold"
-                )
-                repetition_penalty = gr.Slider(
-                    minimum=1.0, maximum=2.0, value=1.1, step=0.05,
-                    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",
-                    info="Maximum length of generated audio (in tokens)"
-                )
+                temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.6, step=0.05, label="Temperature",
+                                        info="Higher values produce more varied speech")
+                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, 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",
+                                               info="Discourage repetition")
+                max_new_tokens = gr.Slider(minimum=100, maximum=2000, value=1200, step=100, label="Max Length",
+                                           info="Maximum generated tokens")
             with gr.Row():
                 submit_btn = gr.Button("Generate Speech", variant="primary")
                 clear_btn = gr.Button("Clear")
@@ -232,7 +235,11 @@ with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
         inputs=[],
         outputs=[text_input, audio_output]
     )
+    
+    gr.Markdown("## ONNX Conversion")
+    onnx_btn = gr.Button("Convert Model to ONNX")
+    onnx_output = gr.Textbox(label="Conversion Output")
+    onnx_btn.click(fn=convert_model_to_onnx, inputs=[], outputs=onnx_output)
 
-# Launch the Gradio app
 if __name__ == "__main__":
     demo.queue().launch(share=False, ssr_mode=False)