| import os as _os |
| HERE = _os.path.dirname(_os.path.abspath(__file__)) |
| ROOT = _os.path.dirname(HERE) |
| import sys, os, json, types, math, shutil; sys.path.insert(0, HERE) |
| import harness, torch |
| from safetensors.torch import save_file |
| pack, nodes = harness.load_pack() |
| shutil.rmtree(harness.OUT, ignore_errors=True); os.makedirs(harness.OUT + "/temp") |
|
|
| |
| import asyncio |
| async def _nodes(): |
| return await (await pack.comfy_entrypoint()).get_node_list() |
| ext = asyncio.run(_nodes()) |
| by_type = {c.__name__: c for c in ext} |
| node_ids = {} |
| for c in ext: |
| sch = c.define_schema(); node_ids[sch.k["node_id"]] = (c, sch) |
| print("node ids:", sorted(node_ids)) |
| WIDGET = {"COMBO", "INT", "FLOAT", "STRING", "BOOLEAN"} |
| for name in ("ogg_text_to_speech", "ogg_voice_clone", "ogg_dialogue", "ogg_voice_design"): |
| w = json.load(open(f"{ROOT}/example_workflows/{name}.json")) |
| for n in w["nodes"]: |
| if n["type"] in ("MarkdownNote", "SaveAudio", "LoadAudio"): continue |
| cls, sch = node_ids[n["type"]] |
| want = [(i.a[0], i.kind in WIDGET) for i in sch.k["inputs"]] |
| have = [(i["name"], "widget" in i) for i in n["inputs"]] |
| assert want == have, (n["type"], want, have) |
| nw = sum(1 for _, wd in want if wd) + sum(1 for i in sch.k["inputs"] if i.k.get("control_after_generate")) |
| if n["type"] != "OggEditor": assert len(n["widgets_values"]) == nw, (n["type"], len(n["widgets_values"]), nw) |
| print("workflow OK:", name) |
|
|
| |
| sys.path.insert(0, HERE) |
| import runpy |
| ns = runpy.run_path(HERE + "/smoke_core.py") |
| tiny, tok = ns["model"], ns["tok"] |
| tmp = harness.OUT + "/tiny_ogg.safetensors" |
| save_file({k: v.contiguous() for k, v in tiny.state_dict().items()}, tmp) |
| nodes.ensure_omnivoice_files = lambda: None |
| seen = {} |
| def fake_from_config_dir(path, dtype=torch.float32, attn="sdpa", prompt_x=False): |
| seen["prompt_x"], seen["path"] = prompt_x, path |
| from transformers import AutoModel |
| m = nodes.of.OmniFish(AutoModel.from_config(tiny.llm.config), prompt_x=prompt_x).to(dtype) |
| t = ns["StubTok"](); m.configure_languages(t); return m, t |
| nodes.of.OmniFish.from_config_dir = staticmethod(fake_from_config_dir) |
| ogg = nodes.load_ogg(tmp, torch.float32) |
| print("load_ogg ok; prompt_x:", seen["prompt_x"], "tokenizer wrapped:", type(ogg.tok).__name__, "| omnivoice dir:", nodes.OMNIVOICE_DIR) |
| try: |
| save_file({"x": torch.zeros(1)}, harness.OUT + "/bad.safetensors"); nodes.load_ogg(harness.OUT + "/bad.safetensors", torch.float32) |
| except ValueError as e: print("rejects a foreign file:", e) |
|
|
| |
| class FakeCodecModel: |
| def encode(self, wav): |
| T = math.ceil(wav.shape[-1] / 2048); torch.manual_seed(1) |
| c = torch.randint(0, 1024, (1, 10, T)); c[:, 0] %= 4096; return c |
| def decode(self, codes, dtype): |
| n = codes.shape[-1] * 2048; t = torch.arange(n) / 44100 |
| return (0.2 * torch.sin(2 * math.pi * 300 * t) * (torch.arange(n) % 20000 < 15000))[None] |
| codec = types.SimpleNamespace(model=FakeCodecModel(), patcher=harness.ModelPatcher(None, torch.device("cpu"), torch.device("cpu"))) |
| audio = {"waveform": torch.randn(1, 1, 48000) * 0.1, "sample_rate": 48000} |
| voice = nodes.OggEncodeVoice.execute(codec, audio, "<|speaker:0|> hello there").args[0] |
| print("encoded voice codes:", tuple(voice["codes"].shape)) |
| nodes.OggSaveVoice.execute(voice, "ogg/voice") |
| vf = [f for f in os.listdir(harness.OUT + "/ogg_voices/ogg")][0] |
| loaded = nodes.OggLoadVoice.execute("ogg/" + vf) if False else None |
| nodes.folder_paths.get_full_path_or_raise = lambda n, f: os.path.join(harness.OUT, n, f) |
| loaded = nodes.OggLoadVoice.execute("ogg/" + vf).args[0] |
| assert torch.equal(loaded["codes"], voice["codes"]) and loaded["transcript"] == voice["transcript"]; print("voice save/load round-trip ok") |
|
|
| r = nodes.OggGenerate.execute(ogg, "[sad] one\ntwo\nthree", "en", 3, 2, 1.5, 0.8, True, 24, "", "", 0.0, 0, 4, loaded).args[0] |
| print("generate: codes", tuple(r["codes"].shape), "takes", r["takes"], "lines", [(l["start"], l["end"]) for l in r["lines"]]) |
| nodes.OggSaveCodes.execute(r, "ogg/take") |
| cf = os.listdir(harness.OUT + "/ogg_codes/ogg")[0] |
| r2 = nodes.OggLoadCodes.execute("ogg/" + cf).args[0] |
| assert torch.equal(r2["codes"], r["codes"]) and r2["takes"] == r["takes"] and r2["lines"] == r["lines"]; print("codes save/load round-trip ok") |
|
|
| a = nodes.OggDecode.execute(codec, r2, False).args[0] |
| expect = sum(r["takes"]) * 2048 + 2 * round(0.3 * 44100) |
| print("decode:", tuple(a["waveform"].shape), a["sample_rate"], "expected samples", expect); assert a["waveform"].shape[-1] == expect |
| a2 = nodes.OggDecode.execute(codec, r2, True).args[0]; print("decode with shorten_long_pauses:", tuple(a2["waveform"].shape)) |
| e = nodes.OggEditor.execute(codec, r2, None, "") |
| print("editor audio:", tuple(e.args[0]["waveform"].shape), "ui:", e.ui["ogg_editor"][0]) |
| assert os.path.exists(os.path.join(harness.OUT, "temp", "ogg_editor", e.ui["ogg_editor"][0]["audio"])) |
| a_ = json.load(open(os.path.join(harness.OUT, "temp", "ogg_editor", e.ui["ogg_editor"][0]["analysis"]))) |
| print("editor analysis: lines", len(a_["lines"]), "speakers", a_["speakers"], "duration", round(a_["duration"], 2)) |
|
|
| |
| from ogg_pack.ogg.editor import DEFAULTS, plan |
| p = plan(a_, dict(DEFAULTS)); print("plan pieces:", len(p["pieces"])) |
|
|
| |
| import importlib; importlib.reload(sys.modules["ogg_pack.nodes"]) if False else None |
| src = open(ROOT + "/nodes.py").read() |
| ns2 = {} |
| print("ALL OK") |
|
|
