| import os as _os |
| HERE = _os.path.dirname(_os.path.abspath(__file__)) |
| ROOT = _os.path.dirname(HERE) |
| import re, sys, time, types, hashlib |
| sys.path.insert(0, ROOT) |
| import torch |
| from transformers import AutoModel, Qwen3Config |
| from ogg import model as of, generate as gen |
|
|
| TEXT_VOCAB = 1000 |
|
|
| class Out: |
| def __init__(self, ids, tensor): self.input_ids = torch.tensor([ids]) if tensor else ids |
|
|
| class StubTok: |
| """stands in for OmniVoice's tokenizer: words and <|special|> markers hashed into the text vocabulary""" |
| def _ids(self, s): |
| parts = re.findall(r"<\|[^|]+\|>|\w+|[^\w\s]", s) |
| return [1 + int(hashlib.md5(p.encode()).hexdigest(), 16) % (TEXT_VOCAB - 2) for p in parts] |
| def __call__(self, s, return_tensors=None, add_special_tokens=True, **kw): |
| return Out(self._ids(s), return_tensors == "pt") |
| def convert_tokens_to_ids(self, t): return self._ids(t)[0] |
|
|
| cfg = Qwen3Config(vocab_size=TEXT_VOCAB, hidden_size=64, intermediate_size=128, num_hidden_layers=2, num_attention_heads=4, |
| num_key_value_heads=2, head_dim=16, max_position_embeddings=2048) |
| cfg._attn_implementation = "sdpa" |
| torch.manual_seed(0) |
| llm = AutoModel.from_config(cfg) |
| model = of.OmniFish(llm).eval() |
| tok = StubTok() |
| model.configure_languages(tok) |
| tok = of.IdTokenizer(tok, model.text_vocab) |
| print("params M:", sum(p.numel() for p in model.parameters()) / 1e6) |
|
|
| |
| state = {k: v.clone() for k, v in model.state_dict().items()} |
| fresh = model.load_weights(state, tok) |
| print("load_weights fresh (False expected, tables present):", fresh) |
|
|
| |
| mx = of.OmniFish(AutoModel.from_config(cfg), prompt_x=True) |
| print("px keys:", sum(".px." in k for k in mx.state_dict())) |
|
|
| prompt = torch.randint(0, 1024, (10, 24)); prompt[0] %= 4096 |
| for label, kw in [("no prompt", dict(prompt_codes=None, prompt_text="")), |
| ("with prompt", dict(prompt_codes=prompt, prompt_text="hello there")), |
| ("left+after ctx", dict(prompt_codes=prompt, prompt_text="hello there", semantic_left=("one", "two"), after=("three",)))]: |
| t = time.time() |
| last = None |
| for clean, codes in gen.blocks_cached(model, tok, text="[sighing] hello world", lang="en", block=16, steps=2, cfg=1.5, max_frames=32, seed=1, **kw): |
| last = (clean, codes) |
| clean, codes = last |
| n = gen.speech_frames(clean) |
| print(f"{label}: clean {tuple(clean.shape)} codes {tuple(codes.shape)} speech {n} range ok:", |
| int(codes[0, :n].max()) < 4096 if n else None, int(codes[1:, :n].max()) < 1024 if n else None, f"{time.time()-t:.1f}s") |
|
|
| |
| for clean, codes in gen.blocks_cached(model, tok, None, "", "hi", "pl", block=16, steps=2, cfg=0, max_frames=16, seed=2): |
| pass |
| print("cfg=0 ok", tuple(codes.shape)) |
|
|
