File size: 8,203 Bytes
0b5b901
87012a8
4189fe1
9bf14d0
87012a8
55515cc
2a24991
d9ea17d
0316ec3
e3958ab
479f253
 
 
2008a3f
1ab029d
e3958ab
 
 
 
 
 
 
 
 
83532d0
f4406f3
e3958ab
 
 
 
479f253
e3958ab
 
 
 
 
 
3d65908
e3958ab
a0cc672
 
 
 
 
e3958ab
 
9bf14d0
0dfc310
9bf14d0
e3958ab
 
9bf14d0
 
e3958ab
5031731
e3958ab
 
0b5b901
 
9bf14d0
5031731
e3958ab
 
bca75ea
d44e840
f63f843
e3958ab
 
 
 
 
 
 
 
 
 
 
 
 
 
0b5b901
7bb84b7
 
 
 
 
e3958ab
9e2fbd8
e3958ab
 
9e2fbd8
e3958ab
0b5b901
 
e3958ab
a8606ac
d44e840
a09ea48
4189fe1
d44e840
e3958ab
 
 
 
 
 
 
a0cc672
f63f843
 
93bffbb
2a24991
 
55515cc
2a24991
55515cc
 
 
2a24991
55515cc
 
 
 
 
 
 
2a24991
55515cc
 
 
 
 
2a24991
55515cc
 
 
 
 
29f8312
55515cc
2a24991
 
 
 
 
 
 
 
 
 
 
e3958ab
 
0238891
 
e3958ab
 
 
 
2a24991
e3958ab
 
 
 
 
0b5b901
e96cc47
 
e3958ab
 
 
a0cc672
 
 
 
 
bca75ea
5031731
479f253
a09ea48
e3958ab
83532d0
 
5031731
479f253
5031731
 
e3958ab
 
 
 
5031731
e3958ab
a4cfefc
e3958ab
83532d0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
# app.py ──────────────────────────────────────────────────────────────
import os, json, torch, asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
from transformers.generation.utils import Cache # Added import
from transformers.cache_utils import DynamicCache # Added import
from snac import SNAC

# 0) Login + Device ---------------------------------------------------
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    login(HF_TOKEN)

device = "cuda" if torch.cuda.is_available() else "cpu"
torch.backends.cuda.enable_flash_sdp(False)          # PyTorch‑2.2‑Bug

# 1) Konstanten -------------------------------------------------------
REPO           = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
CHUNK_TOKENS   = 50
START_TOKEN    = 128259
NEW_BLOCK      = 128257
EOS_TOKEN      = 128258
AUDIO_BASE     = 128266
AUDIO_SPAN   = 4096 * 7                                # 28 672 Codes
AUDIO_IDS  = torch.arange(AUDIO_BASE, AUDIO_BASE + AUDIO_SPAN) # Renamed VALID_AUDIO to AUDIO_IDS

# 2) Logit‑Mask (NEW_BLOCK + Audio; EOS erst nach 1. Block) ----------
class AudioMask(LogitsProcessor):
    def __init__(self, audio_ids: torch.Tensor):
        super().__init__()
        self.allow = torch.cat([
            torch.tensor([NEW_BLOCK], device=audio_ids.device),
            audio_ids
        ])
        self.eos   = torch.tensor([EOS_TOKEN], device=audio_ids.device)
        self.sent_blocks = 0
        self.buffer_pos = 0 # Added buffer position

    def __call__(self, input_ids, scores):
        allow = torch.cat([self.allow, self.eos]) # Reverted masking logic
        mask = torch.full_like(scores, float("-inf"))
        mask[:, allow] = 0
        return scores + mask

# 3) FastAPI Grundgerüst ---------------------------------------------
app = FastAPI()

@app.get("/")
def hello():
    return {"status": "ok"}

@app.on_event("startup")
def load_models():
    global tok, model, snac, masker
    print("⏳ Lade Modelle …", flush=True)

    tok   = AutoTokenizer.from_pretrained(REPO)
    snac  = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
    model = AutoModelForCausalLM.from_pretrained(
        REPO,
        device_map={"": 0} if device == "cuda" else None,
        torch_dtype=torch.bfloat16 if device == "cuda" else None,
        low_cpu_mem_usage=True,
    )
    model.config.pad_token_id = model.config.eos_token_id
    masker = AudioMask(AUDIO_IDS.to(device))

    print("✅ Modelle geladen", flush=True)

# 4) Helper -----------------------------------------------------------
def build_prompt(text: str, voice: str):
    prompt_ids = tok(f"{voice}: {text}", return_tensors="pt").input_ids.to(device)
    ids   = torch.cat([torch.tensor([[START_TOKEN]], device=device),
                       prompt_ids,
                       torch.tensor([[128009, 128260]], device=device)], 1)
    attn  = torch.ones_like(ids)
    return ids, attn

def decode_block(block7: list[int]) -> bytes:
    l1,l2,l3=[],[],[]
    l1.append(block7[0] - 0 * 4096) # Subtract position 0 offset
    l2.append(block7[1] - 1 * 4096) # Subtract position 1 offset
    l3 += [block7[2] - 2 * 4096, block7[3] - 3 * 4096] # Subtract position offsets
    l2.append(block7[4] - 4 * 4096) # Subtract position 4 offset
    l3 += [block7[5] - 5 * 4096, block7[6] - 6 * 4096] # Subtract position offsets

    with torch.no_grad():
        codes = [torch.tensor(x, device=device).unsqueeze(0)
                 for x in (l1,l2,l3)]
        audio = snac.decode(codes).squeeze().detach().cpu().numpy()

    return (audio*32767).astype("int16").tobytes()

# 5) WebSocket‑Endpoint ----------------------------------------------
@app.websocket("/ws/tts")
async def tts(ws: WebSocket):
    await ws.accept()
    try:
        req   = json.loads(await ws.receive_text())
        text  = req.get("text", "")
        voice = req.get("voice", "Jakob")

        ids, attn  = build_prompt(text, voice)
        past       = None
        offset_len = ids.size(1)          # wie viele Tokens existieren schon
        last_tok   = None
        buf         = []

        while True:
            print(f"DEBUG: Before generate - past is None: {past is None}", flush=True) # Added logging
            print(f"DEBUG: Before generate - type of past: {type(past) if past is not None else 'None'}", flush=True) # Added logging)

            if past is None:
                # First generation step
                gen = model.generate(
                    input_ids       = ids,
                    attention_mask  = attn,
                    past_key_values = past, # This will be None
                    max_new_tokens = 1,
                    logits_processor=[masker],
                    do_sample=True, temperature=0.7, top_p=0.95,
                    use_cache=True,
                    return_dict_in_generate=True,
                )
            else:
                # Subsequent generation steps
                current_input_ids = torch.tensor([[last_tok]], device=device)
                current_attention_mask = torch.ones_like(current_input_ids)
                gen = model.generate(
                    input_ids       = current_input_ids,
                    attention_mask  = current_attention_mask,
                    past_key_values = past, # This will be a Cache object
                    max_new_tokens = 1,
                    logits_processor=[masker],
                    do_sample=True, temperature=0.7, top_p=0.95,
                    use_cache=True,
                    return_dict_in_generate=True,
                    cache_position=torch.tensor([offset_len], device=device) # Explicitly pass cache_position
                )

            print(f"DEBUG: After generate - type of gen.past_key_values: {type(gen.past_key_values)}", flush=True) # Added logging)

            # Convert legacy tuple cache to DynamicCache if necessary (only after the first step)
            if past is None and isinstance(gen.past_key_values, tuple):
                past = DynamicCache.from_legacy_cache(gen.past_key_values)
            else:
                # For subsequent steps, just update past with the new cache object
                past = gen.past_key_values

            print(f"DEBUG: After cache update - type of past: {type(past)}", flush=True) # Added logging)

            # ----- neue Tokens heraus schneiden --------------------------
            seq  = gen.sequences[0].tolist()
            new  = seq[offset_len:]
            if not new:                         # nichts -> fertig
                break
            offset_len += len(new)

            # ----- Update last_tok ---------
            last_tok = new[-1]

            print("new tokens:", new[:25], flush=True)

            # ----- Token‑Handling ----------------------------------------
            for t in new:
                if t == EOS_TOKEN: # Re-enabled EOS check
                    raise StopIteration # Re-enabled EOS check
                if t == NEW_BLOCK:
                    buf.clear()
                    continue
                buf.append(t - AUDIO_BASE) # Reverted to appending relative token
                if len(buf) == 7:
                    await ws.send_bytes(decode_block(buf))
                    buf.clear()
                    masker.sent_blocks = 1      # ab jetzt EOS zulässig

    except (StopIteration, WebSocketDisconnect):
        pass
    except Exception as e:
        print("❌ WS‑Error:", e, flush=True)
        import traceback
        traceback.print_exc()
        if ws.client_state.name != "DISCONNECTED":
            await ws.close(code=1011)
    finally:
        if ws.client_state.name != "DISCONNECTED":
            try:
                await ws.close()
            except RuntimeError:
                pass

# 6) Dev‑Start --------------------------------------------------------
if __name__ == "__main__":
    import uvicorn, sys
    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")