Update app.py

app.py

@@ -1,277 +1,116 @@
-import spaces
-from snac import SNAC
+import os
+import json
+import asyncio
 import torch
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-import gradio as gr
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import snapshot_download
 from dotenv import load_dotenv
-load_dotenv()
+from snac import SNAC
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login
 
-# Check if HF_TOKEN is available
-token = os.getenv("HF_TOKEN")
-if token:
-    from huggingface_hub import login
-    login(token=token)
-else:
-    print("⚠️  No HF_TOKEN found – gated model will fail.")
+# — Environment & HF‑Auth —
+load_dotenv()
+HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN:
+    login(token=HF_TOKEN)
 
-# Check if CUDA is available
+# — Device & Modelle laden —
 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)
+snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
 model_name = "canopylabs/3b-de-ft-research_release"
-#"canopylabs/orpheus-3b-0.1-ft"
+print("Loading Orpheus model...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, torch_dtype=torch.bfloat16
+).to(device)
+model.config.pad_token_id = model.config.eos_token_id
 
-# 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.*"
-    ]
-)
-
-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}"
+# — Hilfsfunktionen — 
+def process_prompt(text: str, voice: str):
+    prompt = f"{voice}: {text}"
     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)
+    start = torch.tensor([[128259]], dtype=torch.int64)
+    end = torch.tensor([[128009, 128260]], dtype=torch.int64)
+    ids = torch.cat([start, input_ids, end], dim=1).to(device)
+    mask = torch.ones_like(ids).to(device)
+    return ids, mask
 
-# Parse output tokens to audio
-def parse_output(generated_ids):
+def parse_output(generated_ids: torch.LongTensor):
     token_to_find = 128257
     token_to_remove = 128258
-    
-    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:]
+    idxs = (generated_ids == token_to_find).nonzero(as_tuple=True)[1]
+    if idxs.numel() > 0:
+        last = idxs[-1].item()
+        cropped = generated_ids[:, last+1:]
     else:
-        cropped_tensor = generated_ids
-
-    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
+        cropped = generated_ids
+    # remove padding token markers
+    rows = []
+    for row in cropped:
+        row = row[row != token_to_remove]
+        rows.append(row)
+    flat = rows[0].tolist()
+    # adjust and regroup
+    layer1, layer2, layer3 = [], [], []
+    for i in range(len(flat)//7):
+        base = flat[7*i:7*i+7]
+        layer1.append(base[0])
+        layer2.append(base[1]-4096)
+        layer3.extend([base[2]-(2*4096), base[3]-(3*4096)])
+        layer2.append(base[4]-4*4096)
+        layer3.extend([base[5]-(5*4096), base[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(layer1, device=device).unsqueeze(0),
+        torch.tensor(layer2, device=device).unsqueeze(0),
+        torch.tensor(layer3, 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:
-        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,
-            )
-        
-        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:
-        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 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]
-    )
-
-# Enable queuing for Gradio
-# Create FastAPI app and mount Gradio ASGI app (HTTP mode)
+    audio = snac.decode(codes).detach().squeeze().cpu().numpy()
+    return audio  # float32 numpy at 24000 Hz
 
+# — FastAPI + WebSocket-Endpoint — 
 app = FastAPI()
-app.mount("/", demo.app)
 
-# WebSocket TTS endpoint
 @app.websocket("/ws/tts")
-async def websocket_tts(websocket: WebSocket):
-    await websocket.accept()
+async def tts_ws(ws: WebSocket):
+    await ws.accept()
     try:
         while True:
-            msg = await websocket.receive_text()
+            msg = await ws.receive_text()
             data = json.loads(msg)
             text = data.get("text", "")
-            voice = data.get("voice", VOICES[0])
-            _, audio = generate_speech(text, voice, 0.7, 0.95, 1.1, 1200)
-            chunk_size = 2400  # 0.1s at 24kHz
-            for i in range(0, len(audio), chunk_size):
-                chunk = audio[i:i+chunk_size]
-                await websocket.send_bytes(chunk.tobytes())
-            await websocket.send_text("__END__")
+            voice = data.get("voice", "jana")
+            # Generate tokens
+            ids, mask = process_prompt(text, voice)
+            with torch.no_grad():
+                gen_ids = model.generate(
+                    input_ids=ids,
+                    attention_mask=mask,
+                    max_new_tokens=1200,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.95,
+                    repetition_penalty=1.1,
+                    eos_token_id=128258,
+                )
+            # Convert to waveform
+            audio = parse_output(gen_ids)
+            # PCM16 conversion & chunking
+            pcm16 = (audio * 32767).astype('int16').tobytes()
+            # 0.1 s @24 kHz = 2400 samples = 4800 bytes
+            chunk_size = 2400 * 2
+            for i in range(0, len(pcm16), chunk_size):
+                await ws.send_bytes(pcm16[i:i+chunk_size])
+                await asyncio.sleep(0.1)  # pacing
     except WebSocketDisconnect:
-        print("Client disconnected from /ws/tts")
+        print("Client disconnected")
+    except Exception as e:
+        print("Error in /ws/tts:", e)
+        await ws.close(code=1011)
 
-# Launch when run directly
-def main():
+if __name__ == "__main__":
     import uvicorn
     uvicorn.run("app:app", host="0.0.0.0", port=7860)
-
-if __name__ == "__main__":
-    main()