app.py

# app.py ──────────────────────────────────────────────────────────────
import os
import json
import torch
import asyncio
import traceback # Import traceback for better error logging

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, StoppingCriteria, StoppingCriteriaList
# Import BaseStreamer for the interface
from transformers.generation.streamers import BaseStreamer
from snac import SNAC # Ensure you have 'pip install snac'

# --- Globals (populated in load_models) ---
tok = None
model = None
snac = None
masker = None
stopping_criteria = None
device = "cuda" if torch.cuda.is_available() else "cpu"

# 0) Login + Device ---------------------------------------------------
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    print("🔑 Logging in to Hugging Face Hub...")
    login(HF_TOKEN)

# torch.backends.cuda.enable_flash_sdp(False) # Uncomment if needed for PyTorch‑2.2‑Bug

# 1) Konstanten -------------------------------------------------------
REPO = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
START_TOKEN = 128259
NEW_BLOCK = 128257
EOS_TOKEN = 128258 # Ensure this is correct for the model
AUDIO_BASE = 128266
AUDIO_SPAN = 4096 * 7  # 28672 Codes
CODEBOOK_SIZE = 4096  # Explicitly define the codebook size
AUDIO_IDS_CPU = torch.arange(AUDIO_BASE, AUDIO_BASE + AUDIO_SPAN)

# 2) Logit‑Mask -------------------------------------------------------
class AudioMask(LogitsProcessor):
    def __init__(self, audio_ids: torch.Tensor, new_block_token_id: int, eos_token_id: int):
        super().__init__()
        self.allow = torch.cat([
            torch.tensor([new_block_token_id], device=audio_ids.device, dtype=torch.long),
            audio_ids
        ], dim=0)
        self.eos = torch.tensor([eos_token_id], device=audio_ids.device, dtype=torch.long)
        self.allow_with_eos = torch.cat([self.allow, self.eos], dim=0)
        self.sent_blocks = 0 # State: Number of audio blocks sent

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        current_allow = self.allow_with_eos if self.sent_blocks > 0 else self.allow
        mask = torch.full_like(scores, float("-inf"))
        mask[:, current_allow] = 0
        return scores + mask

    def reset(self):
        self.sent_blocks = 0

# 3) StoppingCriteria für EOS ---------------------------------------
class EosStoppingCriteria(StoppingCriteria):
    def __init__(self, eos_token_id: int):
        self.eos_token_id = eos_token_id

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        if input_ids.shape[1] > 0 and input_ids[:, -1] == self.eos_token_id:
            # print("StoppingCriteria: EOS detected.") # Optional: Uncomment for debugging
            return True
        return False

# 4) Benutzerdefinierter AudioStreamer -------------------------------
class AudioStreamer(BaseStreamer):
    def __init__(self, ws: WebSocket, snac_decoder: SNAC, audio_mask: AudioMask, loop: asyncio.AbstractEventLoop, target_device: str):
        self.ws = ws
        self.snac = snac_decoder
        self.masker = audio_mask
        self.loop = loop
        self.device = target_device
        self.buf: list[int] = []
        self.tasks = set()

    def _decode_block(self, block7: list[int]) -> bytes:
        """
        Decodes a block of 7 audio token values (AUDIO_BASE subtracted) into audio bytes.
        Uses modulo to extract base code value (0-4095).
        Maps extracted values using the structure potentially correct for Kartoffel_Orpheus.
        """
        if len(block7) != 7:
            print(f"Streamer Warning: _decode_block received {len(block7)} tokens, expected 7. Skipping.")
            return b""

        try:
            # --- Extract base code value (0 to CODEBOOK_SIZE-1) for each slot using modulo ---
            code_val_0 = block7[0] % CODEBOOK_SIZE
            code_val_1 = block7[1] % CODEBOOK_SIZE
            code_val_2 = block7[2] % CODEBOOK_SIZE
            code_val_3 = block7[3] % CODEBOOK_SIZE
            code_val_4 = block7[4] % CODEBOOK_SIZE
            code_val_5 = block7[5] % CODEBOOK_SIZE
            code_val_6 = block7[6] % CODEBOOK_SIZE

            # --- Map the extracted code values to the SNAC codebooks (l1, l2, l3) ---
            l1 = [code_val_0]
            l2 = [code_val_1, code_val_4]
            l3 = [code_val_2, code_val_3, code_val_5, code_val_6]

        except IndexError:
            print(f"Streamer Error: Index out of bounds during token mapping. Block: {block7}")
            return b""
        except Exception as e_map:
            print(f"Streamer Error: Exception during code value extraction/mapping: {e_map}. Block: {block7}")
            return b""

        # --- Convert lists to tensors on the correct device ---
        try:
            codes_l1 = torch.tensor(l1, dtype=torch.long, device=self.device).unsqueeze(0)
            codes_l2 = torch.tensor(l2, dtype=torch.long, device=self.device).unsqueeze(0)
            codes_l3 = torch.tensor(l3, dtype=torch.long, device=self.device).unsqueeze(0)
            codes = [codes_l1, codes_l2, codes_l3]
        except Exception as e_tensor:
            print(f"Streamer Error: Exception during tensor conversion: {e_tensor}. l1={l1}, l2={l2}, l3={l3}")
            return b""

        # --- Decode using SNAC ---
        try:
            with torch.no_grad():
                audio = self.snac.decode(codes)[0]
        except Exception as e_decode:
            print(f"Streamer Error: Exception during snac.decode: {e_decode}")
            print(f"Input codes shapes: {[c.shape for c in codes]}")
            print(f"Input codes dtypes: {[c.dtype for c in codes]}")
            print(f"Input codes devices: {[c.device for c in codes]}")
            print(f"Input code values (min/max): L1({min(l1)}/{max(l1)}) L2({min(l2)}/{max(l2)}) L3({min(l3)}/{max(l3)})")
            return b""

        # --- Post-processing ---
        try:
            audio_np = audio.squeeze().detach().cpu().numpy()
            audio_bytes = (audio_np * 32767).astype("int16").tobytes()
            return audio_bytes
        except Exception as e_post:
            print(f"Streamer Error: Exception during post-processing: {e_post}. Audio tensor shape: {audio.shape}")
            return b""

    async def _send_audio_bytes(self, data: bytes):
        """Coroutine to send bytes over WebSocket."""
        if not data:
            return
        try:
            await self.ws.send_bytes(data)
        except WebSocketDisconnect:
            print("Streamer: WebSocket disconnected during send.")
        except Exception as e:
            # Log errors other than expected disconnects more visibly maybe
            if "Cannot call \"send\" once a close message has been sent" not in str(e):
                 print(f"Streamer: Error sending bytes: {e}")
            # else: # Optionally print disconnect errors quietly
            #    print("Streamer: Attempted send after close.")
            pass # Avoid flooding logs if client disconnects early

    def put(self, value: torch.LongTensor):
        """
        Receives new token IDs (Tensor) from generate().
        Processes tokens, decodes full blocks, and schedules sending.
        """
        if value.numel() == 0:
            return
        # Ensure value is on CPU and flatten to a list of ints
        new_token_ids = value.squeeze().cpu().tolist() # Move to CPU before list conversion
        if isinstance(new_token_ids, int):
            new_token_ids = [new_token_ids]

        for t in new_token_ids:
            # --- DEBUGGING PRINT ---
            # Log every token ID received from the model
            print(f"Streamer received token ID: {t}")
            # --- END DEBUGGING ---

            if t == EOS_TOKEN:
                # print("Streamer: EOS token encountered.") # Optional debugging
                break # Stop processing this batch if EOS is found

            if t == NEW_BLOCK:
                # print("Streamer: NEW_BLOCK token encountered.") # Optional debugging
                self.buf.clear()
                continue # Move to the next token

            # Check if token is within the expected audio range
            if AUDIO_BASE <= t < AUDIO_BASE + AUDIO_SPAN:
                self.buf.append(t - AUDIO_BASE) # Store value relative to base
            # else: # Log unexpected tokens if needed
                # print(f"Streamer Warning: Ignoring unexpected token {t} (outside audio range [{AUDIO_BASE}, {AUDIO_BASE + AUDIO_SPAN}))")
                pass

            # If buffer has 7 tokens, decode and send
            if len(self.buf) == 7:
                audio_bytes = self._decode_block(self.buf)
                self.buf.clear() # Clear buffer after processing

                if audio_bytes: # Only send if decoding was successful
                    # Schedule the async send function to run on the main event loop
                    future = asyncio.run_coroutine_threadsafe(self._send_audio_bytes(audio_bytes), self.loop)
                    self.tasks.add(future)
                    # Optional: Remove completed tasks to prevent memory leak if generation is very long
                    future.add_done_callback(self.tasks.discard)

                    # Allow EOS only after the first full block has been processed and scheduled for sending
                    if self.masker.sent_blocks == 0:
                        # print("Streamer: First audio block processed, allowing EOS.")
                        self.masker.sent_blocks = 1 # Update state in the mask

    def end(self):
        """Called by generate() when generation finishes."""
        if len(self.buf) > 0:
            print(f"Streamer: End of generation with incomplete block ({len(self.buf)} tokens). Discarding.")
            self.buf.clear()
        # print(f"Streamer: Generation finished.") # Optional debugging
        pass

# 5) FastAPI App ------------------------------------------------------
app = FastAPI()

@app.on_event("startup")
async def load_models_startup():
    global tok, model, snac, masker, stopping_criteria, device, AUDIO_IDS_CPU, EOS_TOKEN

    print(f"🚀 Starting up on device: {device}")
    print("⏳ Lade Modelle …", flush=True)

    tok = AutoTokenizer.from_pretrained(REPO)
    print("Tokenizer loaded.")

    snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
    print(f"SNAC loaded to {device}.")

    model_dtype = torch.float32
    if device == "cuda":
        if torch.cuda.is_bf16_supported():
            model_dtype = torch.bfloat16
            print("Using bfloat16 for model.")
        else:
            model_dtype = torch.float16
            print("Using float16 for model.")

    model = AutoModelForCausalLM.from_pretrained(
        REPO,
        device_map={"": 0} if device == "cuda" else None,
        torch_dtype=model_dtype,
        low_cpu_mem_usage=True,
    )

    # --- Verify EOS Token ---
    # Use the actual EOS token ID from the loaded model/tokenizer config
    config_eos_id = model.config.eos_token_id
    tokenizer_eos_id = tok.eos_token_id

    if config_eos_id is None:
        print("🚨 WARNING: model.config.eos_token_id is None!")
        # Fallback or default? Let's use the constant for now, but this needs checking.
        final_eos_token_id = EOS_TOKEN
    elif tokenizer_eos_id is not None and config_eos_id != tokenizer_eos_id:
         print(f"⚠️ WARNING: Mismatch! model.config.eos_token_id ({config_eos_id}) != tok.eos_token_id ({tokenizer_eos_id}). Using model config ID.")
         final_eos_token_id = config_eos_id
    else:
         final_eos_token_id = config_eos_id

    # Update the global constant if it differs or wasn't set properly by config
    if final_eos_token_id != EOS_TOKEN:
         print(f"🔄 Updating EOS_TOKEN constant from {EOS_TOKEN} to {final_eos_token_id}")
         EOS_TOKEN = final_eos_token_id # Update the global constant

    # Set pad_token_id to the determined EOS token ID
    model.config.pad_token_id = EOS_TOKEN
    print(f"Using EOS Token ID: {EOS_TOKEN}")
    # --- End Verify EOS Token ---


    print(f"Model loaded to {model.device} with dtype {model.dtype}.")
    model.eval()

    audio_ids_device = AUDIO_IDS_CPU.to(device)
    masker = AudioMask(audio_ids_device, NEW_BLOCK, EOS_TOKEN) # Use updated EOS_TOKEN
    print("AudioMask initialized.")

    stopping_criteria = StoppingCriteriaList([EosStoppingCriteria(EOS_TOKEN)]) # Use updated EOS_TOKEN
    print("StoppingCriteria initialized.")

    print("✅ Modelle geladen und bereit!", flush=True)

@app.get("/")
def hello():
    return {"status": "ok", "message": "TTS Service is running"}

# 6) Helper zum Prompt Bauen -------------------------------------------
def build_prompt(text: str, voice: str) -> tuple[torch.Tensor, torch.Tensor]:
    """Builds the input_ids and attention_mask for the model."""
    prompt_text = f"{voice}: {text}"
    prompt_ids = tok(prompt_text, return_tensors="pt").input_ids.to(device)

    input_ids = torch.cat([
        torch.tensor([[START_TOKEN]], device=device, dtype=torch.long),
        prompt_ids,
        torch.tensor([[NEW_BLOCK]], device=device, dtype=torch.long)
    ], dim=1)

    attention_mask = torch.ones_like(input_ids)
    return input_ids, attention_mask

# 7) WebSocket‑Endpoint (vereinfacht mit Streamer) ---------------------
@app.websocket("/ws/tts")
async def tts(ws: WebSocket):
    await ws.accept()
    print("🔌 Client connected")
    streamer = None
    main_loop = asyncio.get_running_loop()

    try:
        req_text = await ws.receive_text()
        print(f"Received request: {req_text}")
        req = json.loads(req_text)
        text = req.get("text", "Hallo Welt, wie geht es dir heute?")
        voice = req.get("voice", "Jakob")

        if not text:
            print("⚠️ Request text is empty.")
            await ws.close(code=1003, reason="Text cannot be empty")
            return

        print(f"Generating audio for: '{text}' with voice '{voice}'")
        ids, attn = build_prompt(text, voice)
        masker.reset()
        streamer = AudioStreamer(ws, snac, masker, main_loop, device)

        print("Starting generation in background thread...")
        # --- DEBUGGING: Adjusted Generation Parameters ---
        await asyncio.to_thread(
            model.generate,
            input_ids=ids,
            attention_mask=attn,
            max_new_tokens=1500, # Keep lower for faster debugging cycles initially
            logits_processor=[masker],
            stopping_criteria=stopping_criteria,
            # --- Adjusted Parameters for Debugging Repetition ---
            do_sample=True,
            temperature=0.7,     # Slightly higher temperature
            # top_p=0.9,         # Commented out top_p for simpler testing
            repetition_penalty=1.2, # Slightly stronger penalty
            # --- End Adjusted Parameters ---
            use_cache=True,
            streamer=streamer
        )
        print("Generation thread finished.")

    except WebSocketDisconnect:
        print("🔌 Client disconnected.")
    except json.JSONDecodeError:
        print("❌ Invalid JSON received.")
        if ws.client_state.name == "CONNECTED":
            await ws.close(code=1003, reason="Invalid JSON format")
    except Exception as e:
        error_details = traceback.format_exc()
        print(f"❌ WS‑Error: {e}\n{error_details}", flush=True)
        error_payload = json.dumps({"error": str(e)})
        try:
            if ws.client_state.name == "CONNECTED":
                 await ws.send_text(error_payload)
        except Exception:
            pass
        if ws.client_state.name == "CONNECTED":
            await ws.close(code=1011)
    finally:
        if streamer:
            try:
                streamer.end()
            except Exception as e_end:
                 print(f"Error during streamer.end(): {e_end}")

        print("Closing connection.")
        if ws.client_state.name == "CONNECTED":
            try:
                await ws.close(code=1000)
            except RuntimeError as e_close:
                 print(f"Runtime error closing websocket: {e_close}")
            except Exception as e_close_final:
                 print(f"Error closing websocket: {e_close_final}")
        elif ws.client_state.name != "DISCONNECTED":
             print(f"WebSocket final state: {ws.client_state.name}")
        print("Connection closed.")

# 8) Dev‑Start --------------------------------------------------------
if __name__ == "__main__":
    import uvicorn
    print("Starting Uvicorn server...")
    # Note: Consider running with --workers 1 if you face issues with globals/GPU memory
    # uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info", workers=1)
    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")