Update Maya1 Gradio app with preset characters

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
-import asyncio
+import torch
 import io
-import sys
-sys.path.insert(0, '.')
+import wave
+import numpy as np
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from snac import SNAC
 
 # Mock spaces module for local testing
 try:
@@ -14,11 +16,18 @@ except ImportError:
             return func
     spaces = SpacesMock()
 
-from maya1.model_loader import Maya1Model
-from maya1.pipeline import Maya1Pipeline
-from maya1.prompt_builder import Maya1PromptBuilder
-from maya1.snac_decoder import SNACDecoder
-from maya1.constants import AUDIO_SAMPLE_RATE
+# Constants
+CODE_START_TOKEN_ID = 128257
+CODE_END_TOKEN_ID = 128258
+CODE_TOKEN_OFFSET = 128266
+SNAC_MIN_ID = 128266
+SNAC_MAX_ID = 156937
+SOH_ID = 128259
+EOH_ID = 128260
+SOA_ID = 128261
+BOS_ID = 128000
+TEXT_EOT_ID = 128009
+AUDIO_SAMPLE_RATE = 24000
 
 # Preset characters (2 realistic + 2 creative)
 PRESET_CHARACTERS = {
@@ -40,53 +49,77 @@ PRESET_CHARACTERS = {
     }
 }
 
-# Global pipeline variables
+# Global model variables
 model = None
-prompt_builder = None
-snac_decoder = None
-pipeline = None
+tokenizer = None
+snac_model = None
 models_loaded = False
 
-def load_models():
-    """Load Maya1 vLLM model and pipeline (runs once)."""
-    global model, prompt_builder, snac_decoder, pipeline, models_loaded
+def build_prompt(tokenizer, description: str, text: str) -> str:
+    """Build formatted prompt for Maya1."""
+    soh_token = tokenizer.decode([SOH_ID])
+    eoh_token = tokenizer.decode([EOH_ID])
+    soa_token = tokenizer.decode([SOA_ID])
+    sos_token = tokenizer.decode([CODE_START_TOKEN_ID])
+    eot_token = tokenizer.decode([TEXT_EOT_ID])
+    bos_token = tokenizer.bos_token
     
-    if models_loaded:
-        return
+    formatted_text = f'<description="{description}"> {text}'
+    prompt = (
+        soh_token + bos_token + formatted_text + eot_token +
+        eoh_token + soa_token + sos_token
+    )
+    return prompt
+
+def unpack_snac_from_7(snac_tokens: list) -> list:
+    """Unpack 7-token SNAC frames to 3 hierarchical levels."""
+    if snac_tokens and snac_tokens[-1] == CODE_END_TOKEN_ID:
+        snac_tokens = snac_tokens[:-1]
     
-    import torch
-    import os
+    frames = len(snac_tokens) // 7
+    snac_tokens = snac_tokens[:frames * 7]
     
-    # Ensure CUDA is available for HF Spaces
-    if not torch.cuda.is_available():
-        print("Warning: CUDA not available, using CPU")
-        device = "cpu"
-    else:
-        device = "cuda"
-        print(f"CUDA available: {torch.cuda.get_device_name(0)}")
+    if frames == 0:
+        return [[], [], []]
     
-    # Set environment variable for vLLM
-    os.environ.setdefault("VLLM_USE_V1", "0")
+    l1, l2, l3 = [], [], []
     
-    print("Loading Maya1 model with vLLM...")
-    model = Maya1Model(
-        model_path="maya-research/maya1",
-        dtype="bfloat16",
-        max_model_len=8192,
-        gpu_memory_utilization=0.85,
-    )
+    for i in range(frames):
+        slots = snac_tokens[i*7:(i+1)*7]
+        l1.append((slots[0] - CODE_TOKEN_OFFSET) % 4096)
+        l2.extend([
+            (slots[1] - CODE_TOKEN_OFFSET) % 4096,
+            (slots[4] - CODE_TOKEN_OFFSET) % 4096,
+        ])
+        l3.extend([
+            (slots[2] - CODE_TOKEN_OFFSET) % 4096,
+            (slots[3] - CODE_TOKEN_OFFSET) % 4096,
+            (slots[5] - CODE_TOKEN_OFFSET) % 4096,
+            (slots[6] - CODE_TOKEN_OFFSET) % 4096,
+        ])
     
-    print("Initializing prompt builder...")
-    prompt_builder = Maya1PromptBuilder(model.tokenizer, model)
+    return [l1, l2, l3]
+
+def load_models():
+    """Load Maya1 Transformers model (runs once)."""
+    global model, tokenizer, snac_model, models_loaded
     
-    print("Loading SNAC decoder...")
-    snac_decoder = SNACDecoder(
-        device=device,
-        enable_batching=False,
+    if models_loaded:
+        return
+    
+    print("Loading Maya1 model with Transformers...")
+    model = AutoModelForCausalLM.from_pretrained(
+        "maya-research/maya1", 
+        torch_dtype=torch.bfloat16, 
+        device_map="auto",
+        trust_remote_code=True
     )
+    tokenizer = AutoTokenizer.from_pretrained("maya-research/maya1", trust_remote_code=True)
     
-    print("Initializing pipeline...")
-    pipeline = Maya1Pipeline(model, prompt_builder, snac_decoder)
+    print("Loading SNAC decoder...")
+    snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval()
+    if torch.cuda.is_available():
+        snac_model = snac_model.to("cuda")
     
     models_loaded = True
     print("Models loaded successfully!")
@@ -100,7 +133,7 @@ def preset_selected(preset_name):
 
 @spaces.GPU
 def generate_speech(preset_name, description, text, temperature, max_tokens):
-    """Generate emotional speech from description and text using vLLM."""
+    """Generate emotional speech from description and text using Transformers."""
     try:
         # Load models if not already loaded
         load_models()
@@ -115,42 +148,65 @@ def generate_speech(preset_name, description, text, temperature, max_tokens):
         
         print(f"Generating with temperature={temperature}, max_tokens={max_tokens}...")
         
-        # Generate audio using vLLM pipeline (async wrapper)
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        audio_bytes = loop.run_until_complete(
-            pipeline.generate_speech(
-                description=description,
-                text=text,
-                temperature=temperature,
-                top_p=0.9,
-                max_tokens=max_tokens,
+        # Build prompt
+        prompt = build_prompt(tokenizer, description, text)
+        inputs = tokenizer(prompt, return_tensors="pt")
+        
+        if torch.cuda.is_available():
+            inputs = {k: v.to("cuda") for k, v in inputs.items()}
+        
+        # Generate tokens
+        with torch.inference_mode():
+            outputs = model.generate(
+                **inputs, 
+                max_new_tokens=max_tokens,
+                min_new_tokens=28,
+                temperature=temperature, 
+                top_p=0.9, 
                 repetition_penalty=1.1,
-                seed=None,
+                do_sample=True,
+                eos_token_id=CODE_END_TOKEN_ID,
+                pad_token_id=tokenizer.pad_token_id,
             )
-        )
-        loop.close()
         
-        if audio_bytes is None:
-            return None, "Error: Audio generation failed. Try different text or increase max_tokens."
+        # Extract SNAC tokens
+        generated_ids = outputs[0, inputs['input_ids'].shape[1]:].tolist()
+        
+        # Find EOS and extract SNAC codes
+        eos_idx = generated_ids.index(CODE_END_TOKEN_ID) if CODE_END_TOKEN_ID in generated_ids else len(generated_ids)
+        snac_tokens = [t for t in generated_ids[:eos_idx] if SNAC_MIN_ID <= t <= SNAC_MAX_ID]
+        
+        if len(snac_tokens) < 7:
+            return None, "Error: Not enough tokens generated. Try different text or increase max_tokens."
+        
+        # Unpack and decode
+        levels = unpack_snac_from_7(snac_tokens)
+        frames = len(levels[0])
         
-        # Convert bytes to WAV file
-        import wave
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        codes_tensor = [torch.tensor(level, dtype=torch.long, device=device).unsqueeze(0) for level in levels]
+        
+        with torch.inference_mode():
+            z_q = snac_model.quantizer.from_codes(codes_tensor)
+            audio = snac_model.decoder(z_q)[0, 0].cpu().numpy()
+        
+        # Trim warmup
+        if len(audio) > 2048:
+            audio = audio[2048:]
+        
+        # Convert to WAV
+        audio_int16 = (audio * 32767).astype(np.int16)
         wav_buffer = io.BytesIO()
         with wave.open(wav_buffer, 'wb') as wav_file:
             wav_file.setnchannels(1)
             wav_file.setsampwidth(2)
             wav_file.setframerate(AUDIO_SAMPLE_RATE)
-            wav_file.writeframes(audio_bytes)
+            wav_file.writeframes(audio_int16.tobytes())
         
         wav_buffer.seek(0)
-        
-        # Calculate duration
-        duration = len(audio_bytes) // 2 / AUDIO_SAMPLE_RATE
-        frames = len(audio_bytes) // 2 // (AUDIO_SAMPLE_RATE // 6.86) // 7
+        duration = len(audio) / AUDIO_SAMPLE_RATE
         
         status_msg = f"Generated {duration:.2f}s of emotional speech!"
-        
         return wav_buffer, status_msg
     
     except Exception as e:

requirements.txt

@@ -1,10 +1,8 @@
 torch>=2.5.0
 transformers>=4.57.0
 gradio>=5.0.0
-vllm>=0.11.0
 snac>=1.2.1
 soundfile>=0.13.0
 numpy>=2.1.0
 accelerate>=1.10.0
-xformers>=0.0.32