| import os, sys; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import os, math, types, runpy |
| import harness, torch |
| pack, nodes = harness.load_pack() |
| HERE_ = os.path.dirname(os.path.abspath(__file__)) |
|
|
| |
| P = nodes.parse_script |
| names = ["Jean", "Maria", "", "", "", ""] |
| assert P("Jean: hi\nMaria: [sad] ok: fine\nand more\njean: bye", names) == [(0, "hi"), (1, "[sad] ok: fine"), (1, "and more"), (0, "bye")] |
| assert P("Maria:\nline one\nline two", names) == [(1, "line one"), (1, "line two")] |
| assert P("Жан: Привет\nМария: Да", ["Жан", "Мария"] + [""] * 4) == [(0, "Привет"), (1, "Да")] |
| for bad, why in [("hi there", "no name first"), ("Jean: a\nMaía: b", "misspelled speaker"), ("Jean: a\nHe said: yes", "unknown plain word before colon")]: |
| try: P(bad, names); raise SystemExit("no error for: " + why) |
| except ValueError as e: print("rejects", why, "->", str(e)[:70]) |
| assert P("Jean: a\nthe time is 5:30, ok", names) == [(0, "a"), (0, "the time is 5:30, ok")] |
| print("parse_script ok") |
|
|
| |
| ns = runpy.run_path(os.path.join(HERE_, "smoke_core.py")) |
| ogg = nodes.Ogg(ns["model"], ns["tok"], types.SimpleNamespace()) |
| vA = {"transcript": "<|speaker:0|> hello there", "codes": torch.randint(0, 1024, (10, 20))} |
| vB = {"transcript": "good morning to you", "codes": torch.randint(0, 1024, (10, 30))} |
| script = "Jean: [excited] one\nMaria: [sighing] two\nJean: three" |
| kw = dict(name_1="Jean", name_2="Maria", language="ru", seed=5, steps=2, cfg=1.5, temperature=0.8, max_frames=24) |
|
|
| r = nodes.OggDialogue.execute(ogg, script, mode="line by line", voice_1=vA, voice_2=vB, **kw).args[0] |
| print("line by line:", tuple(r["codes"].shape), r["takes"], [(l["speaker"], l["text"]) for l in r["lines"]]) |
| assert [l["speaker"] for l in r["lines"]] == [0, 1, 0] and sum(r["takes"]) == r["codes"].shape[1] |
|
|
| r = nodes.OggDialogue.execute(ogg, script, mode="one pass", voice_1=vA, voice_2=vB, **kw).args[0] |
| print("one pass:", tuple(r["codes"].shape), r["takes"], len(r["lines"])) |
| assert len(r["takes"]) == 1 and r["codes"].shape[0] == 10 |
|
|
| |
| seen = {} |
| orig = nodes.gen.blocks_cached |
| def spy(*a, **k): |
| seen.update(prompt_codes=None if a[2] is None else tuple(a[2].shape), prompt_text=a[3], text=a[4], speaker=k.get("speaker"), |
| voices=[(v[0], v[1], tuple(v[2].shape)) for v in k.get("voices") or []]) |
| return orig(*a, **k) |
| nodes.gen.blocks_cached = spy |
| nodes.OggDialogue.execute(ogg, "Jean: a\nJean: b\nMaria: c\nJean: d", mode="one pass", voice_1=vA, voice_2=vB, **kw) |
| print("as trained:", seen) |
| assert seen["text"] == "<|quad_start|>S1<|quad_end|>a b <|quad_start|>S2<|quad_end|>c <|quad_start|>S1<|quad_end|>d" |
| assert seen["speaker"] == "dialogue" and seen["voices"] == [] and seen["prompt_codes"] == (10, 50) |
| assert seen["prompt_text"] == "<|quad_start|>S1<|quad_end|>hello there <|quad_start|>S2<|quad_end|>good morning to you" |
| nodes.OggDialogue.execute(ogg, "Maria: c\nJean: a", mode="one pass", dialogue_layout="reference script", voice_1=vA, voice_2=vB, **kw) |
| print("reference script:", seen) |
| assert seen["text"] == "<|quad_start|>S1<|quad_end|>c <|quad_start|>S2<|quad_end|>a" and seen["prompt_codes"] is None |
| assert seen["voices"] == [("S1", "good morning to you", (10, 30)), ("S2", "hello there", (10, 20))] |
| nodes.gen.blocks_cached = orig |
|
|
| |
| seen = {} |
| orig = nodes.gen.blocks_cached |
| def spy(*a, **k): seen.setdefault("calls", []).append((a[2] is None, k.get("speaker"))); return orig(*a, **k) |
| nodes.gen.blocks_cached = spy |
| r = nodes.OggDialogue.execute(ogg, script, mode="line by line", voice_1=vA, **kw).args[0] |
| nodes.gen.blocks_cached = orig |
| print("no voice for Maria:", seen["calls"]) |
| assert [c[0] for c in seen["calls"]] == [False, True, False] and seen["calls"][1][1] == nodes.resolve_speaker("", 0, "Maria") and seen["calls"][0][1] is None |
| assert r["takes"] == [32, 32, 32] |
| try: nodes.OggDialogue.execute(ogg, script, mode="one pass", voice_1=vA, **kw); raise SystemExit("no error") |
| except ValueError as e: print("one pass without a voice ->", str(e)[:90]) |
| |
| nodes.gen.blocks_cached = spy; seen.clear() |
| nodes.OggDialogue.execute(ogg, "Maria: x", mode="line by line", speaker_2="ru_x", **kw) |
| nodes.gen.blocks_cached = orig; assert seen["calls"] == [(True, "ru_x")] |
| try: nodes.OggDialogue.execute(ogg, "Maria: x", mode="line by line", speaker_2="pl_3", **kw); raise SystemExit("no error") |
| except ValueError as e: print("bad token ->", str(e)[:70]) |
| |
| r = nodes.OggDialogue.execute(ogg, "Jean: only me", mode="line by line", voice_1=vA, **kw).args[0] |
| print("unused speaker needs no voice ok:", r["takes"]) |
|
|
| |
| R = nodes.resolve_speaker |
| assert R("") is None and R(" ") is None |
| assert R("", 1, "Maria") == R("", 99, " maria ") and R("", 1, "Maria") != R("", 1, "Jean") |
| assert R("random", 7) == R("random", 7) and R("random", 7) != R("random", 8) and R("random", 7).startswith("spk_") |
| assert R("spk_5") == "spk_5" and R("ru_x") == "ru_x" |
| for bad in ("pl_3", "q5", "spk_99999", "foo"): |
| try: R(bad); raise SystemExit("no error for " + bad) |
| except ValueError: pass |
| assert nodes.check_quality("q35") == "q35" and nodes.check_quality("") is None |
| try: nodes.check_quality("spk_1"); raise SystemExit("no error") |
| except ValueError: pass |
| print("speaker and quality tokens ok:", R("random", 7), R("", 0, "Jean")) |
|
|
| |
| nodes.gen.blocks_cached = spy; seen.clear() |
| nodes.OggGenerate.execute(ogg, "[sad] hi", "en", 4, 2, 1.5, 0.8, True, 16, "random", "q7", 0.0, 0, 4, None) |
| nodes.gen.blocks_cached = orig; assert seen["calls"] == [(True, R("random", 4))], seen |
| print("generate with random speaker ok:", seen["calls"]) |
|
|
| take = {"codes": torch.arange(10 * 60).reshape(10, 60) % 1000, "lines": [{"speaker": 0, "text": "[sad] hello [pause] there", "chunk": 0, "start": 0, "end": 25}, |
| {"speaker": 0, "text": "second", "chunk": 1, "start": 25, "end": 60}], "takes": [25, 35]} |
| v = nodes.OggVoiceFromTake.execute(take, 0, True).args[0] |
| assert v["transcript"] == "hello there" and tuple(v["codes"].shape) == (10, 25) and torch.equal(v["codes"], take["codes"][:, :25]) |
| v = nodes.OggVoiceFromTake.execute(take, 1, False).args[0]; assert v["transcript"] == "second" and v["codes"].shape[1] == 35 |
| assert nodes.OggVoiceFromTake.execute(take, 0, False).args[0]["transcript"] == "[sad] hello [pause] there" |
| try: nodes.OggVoiceFromTake.execute(take, 2, True); raise SystemExit("no error") |
| except ValueError as e: print("take out of range ->", e) |
| print("voice from take ok") |
| print("DIALOGUE OK") |
|
|
