Update app.py

app.py

@@ -1,113 +1,246 @@
-from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
-from fastapi.responses import StreamingResponse, JSONResponse
-import outetts
-import io
-import json
-import base64
-import struct
-import os 
-# Initialize the interface
-interface = outetts.Interface(
-    config=outetts.ModelConfig.auto_config(
-        model=outetts.Models.VERSION_1_0_SIZE_1B,
-        # For llama.cpp backend
-        #backend=outetts.Backend.LLAMACPP,
-        #quantization=outetts.LlamaCppQuantization.FP16
-        # For transformers backend
-        backend=outetts.Backend.HF,
-    )
+import spaces
+from snac import SNAC
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+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"
+
+print("Loading SNAC model...")
+snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
+snac_model = snac_model.to(device)
+
+model_name = "canopylabs/3b-de-ft-research_release"
+#"canopylabs/orpheus-3b-0.1-ft"
+
+# Download only model config and safetensors
+snapshot_download(
+    repo_id=model_name,
+    allow_patterns=[
+        "config.json",
+        "*.safetensors",
+        "model.safetensors.index.json",
+    ],
+    ignore_patterns=[
+        "optimizer.pt",
+        "pytorch_model.bin",
+        "training_args.bin",
+        "scheduler.pt",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "special_tokens_map.json",
+        "vocab.json",
+        "merges.txt",
+        "tokenizer.*"
+    ]
 )
 
-# Load the default speaker profile
-speaker = interface.load_default_speaker("EN-FEMALE-1-NEUTRAL")
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+model.to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+print(f"Orpheus model loaded to {device}")
+
+# Process text prompt
+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
+    attention_mask = torch.ones_like(modified_input_ids)
+    
+    return modified_input_ids.to(device), attention_mask.to(device)
 
-app = FastAPI()
+# Parse output tokens to audio
+def parse_output(generated_ids):
+    token_to_find = 128257
+    token_to_remove = 128258
+    
+    token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+    if len(token_indices[1]) > 0:
+        last_occurrence_idx = token_indices[1][-1].item()
+        cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
+    else:
+        cropped_tensor = generated_ids
 
-@app.websocket("/ws/tts")
-async def websocket_tts(websocket: WebSocket):
-    await websocket.accept()
+    processed_rows = []
+    for row in cropped_tensor:
+        masked_row = row[row != token_to_remove]
+        processed_rows.append(masked_row)
+
+    code_lists = []
+    for row in processed_rows:
+        row_length = row.size(0)
+        new_length = (row_length // 7) * 7
+        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
+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
+    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)
+    ]
+    
+    audio_hat = snac_model.decode(codes)
+    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
+    
     try:
-        while True:
-            # Empfange Text-Chunk vom Client
-            data = await websocket.receive_text()
-            # Status: Warming up
-            await websocket.send_text(json.dumps({"generation_status": "Warming up TTS model"}))
-            output = interface.generate(
-                config=outetts.GenerationConfig(
-                    text=data,
-                    generation_type=outetts.GenerationType.CHUNKED,
-                    speaker=speaker,
-                    sampler_config=outetts.SamplerConfig(
-                        temperature=0.4
-                    ),
-                )
+        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():
+            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,
             )
-            # Status: Generating linguistic features
-            await websocket.send_text(json.dumps({"generation_status": "Generating linguistic features"}))
-            # Save to buffer
-            import uuid
-            temp_path = f"temp_{uuid.uuid4().hex}.wav"
-            output.save(temp_path)
-            chunk_size = 4096
-            try:
-                with open(temp_path, "rb") as f:
-                    wav_data = f.read()
-                    # WAV header is typically 44 bytes, but let's detect it robustly
-                    # Find the end of the header (data chunk)
-                    if wav_data[:4] != b'RIFF' or wav_data[8:12] != b'WAVE':
-                        raise ValueError("Not a valid WAV file")
-                    # Find 'data' subchunk
-                    data_offset = wav_data.find(b'data')
-                    if data_offset == -1:
-                        raise ValueError("No 'data' chunk found in WAV file")
-                    header_end = data_offset + 8  # 'data' + size (4 bytes)
-                    wav_header = bytearray(wav_data[:header_end])
-                    pcm_data = wav_data[header_end:]
-                    # Patch header: set data length to 0xFFFFFFFF (unknown/streaming)
-                    wav_header[data_offset+4:data_offset+8] = (0xFFFFFFFF).to_bytes(4, 'little')
-                    # Send header + first PCM chunk
-                    first_chunk = pcm_data[:chunk_size]
-                    audio_b64 = base64.b64encode(wav_header + first_chunk).decode("ascii")
-                    await websocket.send_text(json.dumps({
-                        "data": {
-                            "audio_bytes": audio_b64,
-                            "duration": None,
-                            "request_finished": False
-                        }
-                    }))
-                    # Send rest of PCM data in chunks (without header)
-                    idx = chunk_size
-                    while idx < len(pcm_data):
-                        chunk = pcm_data[idx:idx+chunk_size]
-                        if not chunk:
-                            break
-                        audio_b64 = base64.b64encode(chunk).decode("ascii")
-                        await websocket.send_text(json.dumps({
-                            "data": {
-                                "audio_bytes": audio_b64,
-                                "duration": None,
-                                "request_finished": False
-                            }
-                        }))
-                        idx += chunk_size
-            finally:
-                try:
-                    os.remove(temp_path)
-                except FileNotFoundError:
-                    pass
-            # Final event
-            await websocket.send_text(json.dumps({
-                "data": {
-                    "audio_bytes": "",
-                    "duration": None,
-                    "request_finished": True
-                }
-            }))
-    except WebSocketDisconnect:
-        pass
+        
+        progress(0.6, "Processing speech tokens...")
+        code_list = parse_output(generated_ids)
+        
+        progress(0.8, "Converting to audio...")
+        audio_samples = redistribute_codes(code_list, snac_model)
+        
+        return (24000, audio_samples)  # Return sample rate and audio
     except Exception as e:
-        await websocket.send_text(json.dumps({"error": str(e)}))
+        print(f"Error generating speech: {e}")
+        return None
+
+# Examples for the 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 <sigh> like sighing, or <laugh> laughing, or <yawn> yawning!", "dan", 0.7, 0.95, 1.1, 1200],
+    ["I live in San Francisco, and have, uhm let's see, 3 billion 7 hundred ... <gasp> well, lets just say a lot of parameters.", "leah", 0.6, 0.9, 1.2, 1200],
+    ["Sometimes when I talk too much, I need to <cough> excuse myself. <sniffle> The weather has been quite cold lately.", "leo", 0.65, 0.9, 1.1, 1200],
+    ["Public speaking can be challenging. <groan> But with enough practice, anyone can become better at it.", "jess", 0.7, 0.95, 1.1, 1200],
+    ["The hike was exhausting but the view from the top was absolutely breathtaking! <sigh> It was totally worth it.", "mia", 0.65, 0.9, 1.15, 1200],
+    ["Did you hear that joke? <laugh> I couldn't stop laughing when I first heard it. <chuckle> It's still funny.", "zac", 0.7, 0.95, 1.1, 1200],
+    ["After running the marathon, I was so tired <yawn> and needed a long rest. <sigh> But I felt accomplished.", "zoe", 0.6, 0.95, 1.1, 1200]
+]
+
+# Available voices
+VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]
+
+# Available Emotive Tags
+EMOTIVE_TAGS = ["`<laugh>`", "`<chuckle>`", "`<sigh>`", "`<cough>`", "`<sniffle>`", "`<groan>`", "`<yawn>`", "`<gasp>`"]
+
+# Create Gradio interface
+with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
+    gr.Markdown(f"""
+    # 🎵 [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.
+    
+    ## Tips for better prompts:
+    - Add paralinguistic elements like {", ".join(EMOTIVE_TAGS)} or `uhm` for more human-like speech.
+    - Longer text prompts generally work better than very short phrases
+    - Increasing `repetition_penalty` and `temperature` makes the model speak faster.
+    """)    
+    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"
+            )
+            
+            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)"
+                )
+            
+            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],
+        outputs=audio_output,
+        fn=generate_speech,
+        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
+if __name__ == "__main__":
+    demo.queue().launch(share=False, ssr_mode=False)