Spaces:
Running on Zero
Running on Zero
AndrianBalanescu
fix(burst): route UNET through ModelSamplingSD3 shift=8 per official Comfy-Org Wan2.1 template
c03e859 Download tests/unit/test_hyperstack_burst.py from abalanescu/flow2: direct link, hf CLI and curl.
- Browser
- Download file 8.3 kB
-
https://huggingface.co/spaces/abalanescu/flow2/resolve/main/tests/unit/test_hyperstack_burst.py
- Command line
-
hf download hf://spaces/abalanescu/flow2/tests/unit/test_hyperstack_burst.py
-
curl -L -o test_hyperstack_burst.py https://huggingface.co/spaces/abalanescu/flow2/resolve/main/tests/unit/test_hyperstack_burst.py
8.3 kB
| """Offline unit tests for hyperstack_burst.py. | |
| No network, no VM creation, no GPU: manifest parsing/validation, workflow | |
| filling, frame math, cost estimation, and key sourcing are tested in isolation | |
| by importing the module (hyperstack SDK import is deferred inside functions). | |
| """ | |
| import json | |
| import re | |
| import sys | |
| from pathlib import Path | |
| import pytest | |
| SCRIPTS = Path(__file__).resolve().parents[2] / "moldovan-qwen" / "scripts" | |
| sys.path.insert(0, str(SCRIPTS)) | |
| import hyperstack_burst as hb # noqa: E402 | |
| def artwork_dir(tmp_path): | |
| d = tmp_path / "artwork" | |
| d.mkdir() | |
| for i in (1, 2, 3): | |
| (d / f"cover_{i}.png").write_bytes(b"\x89PNG\r\n\x1a\n" + b"stub" * 8) | |
| return d | |
| def scenes(artwork_dir): | |
| return [ | |
| {"scene_num": 1, "artwork": "cover_1.png", "prompt": "zoom", "duration": 5, "resolution": "480p"}, | |
| {"scene_num": 2, "artwork": "cover_2.png", "prompt": "pan", "duration": 4.5, "resolution": "720p"}, | |
| ] | |
| def test_load_manifest_list(scenes, tmp_path): | |
| p = tmp_path / "m.json" | |
| p.write_text(json.dumps(scenes)) | |
| assert hb.load_manifest(p) == scenes | |
| def test_load_manifest_dict_wrapper(scenes, tmp_path): | |
| p = tmp_path / "m.json" | |
| p.write_text(json.dumps({"scenes": scenes})) | |
| assert hb.load_manifest(p) == scenes | |
| def test_load_manifest_rejects_empty(tmp_path): | |
| p = tmp_path / "m.json" | |
| p.write_text(json.dumps([])) | |
| with pytest.raises(ValueError): | |
| hb.load_manifest(p) | |
| def test_validate_scenes_ok(scenes, artwork_dir): | |
| assert hb.validate_scenes(scenes, artwork_dir) == [] | |
| def test_validate_scenes_missing_artwork(scenes, artwork_dir): | |
| scenes[0]["artwork"] = "nope.png" | |
| errors = hb.validate_scenes(scenes, artwork_dir) | |
| assert any("cover" not in e and "missing" in e for e in errors) | |
| def test_validate_scenes_missing_prompt(artwork_dir): | |
| scenes = [{"scene_num": 1, "artwork": "cover_1.png"}] | |
| errors = hb.validate_scenes(scenes, artwork_dir) | |
| assert any("prompt" in e for e in errors) | |
| def test_validate_scenes_bad_resolution(scenes, artwork_dir): | |
| scenes[0]["resolution"] = "1080p" | |
| assert any("resolution" in e for e in hb.validate_scenes(scenes, artwork_dir)) | |
| def test_validate_scenes_bad_duration(scenes, artwork_dir): | |
| scenes[0]["duration"] = 30 | |
| assert any("duration" in e for e in hb.validate_scenes(scenes, artwork_dir)) | |
| def test_frames_for_is_4n_plus_1(): | |
| assert hb.frames_for(5) == 81 # 5s * 16fps / 4 -> 20 -> 4*20+1 | |
| for d in (1, 2, 3, 4.5, 5, 7.5, 10): | |
| f = hb.frames_for(d) | |
| assert f >= 5 and (f - 1) % 4 == 0 | |
| def test_fill_workflow_injects_scene_fields(scenes, artwork_dir): | |
| wf = hb.fill_workflow(scenes[0], artwork_dir) | |
| assert wf["42"]["inputs"]["text"] == "zoom" | |
| assert wf["44"]["inputs"]["image"] == "cover_1.png" | |
| assert wf["45"]["inputs"]["length"] == 81 | |
| assert wf["45"]["inputs"]["width"] == 832 | |
| assert wf["46"]["inputs"]["seed"] == 1001 | |
| assert wf["49"]["inputs"]["filename_prefix"] == "burst-scene_001" | |
| # negative default present | |
| assert "blurry" in wf["43"]["inputs"]["text"] | |
| def test_fill_workflow_720p_uses_720_unet(scenes, artwork_dir): | |
| wf = hb.fill_workflow(scenes[1], artwork_dir) | |
| assert "720p" in wf["37"]["inputs"]["unet_name"] | |
| assert wf["45"]["inputs"]["width"] == 1280 | |
| def test_fill_workflow_template_not_mutated(scenes, artwork_dir): | |
| before = json.dumps(hb.WORKFLOW_TEMPLATE) | |
| hb.fill_workflow(scenes[0], artwork_dir) | |
| assert json.dumps(hb.WORKFLOW_TEMPLATE) == before | |
| def test_workflow_graph_integrity(scenes, artwork_dir): | |
| """All internal references point to existing nodes with correct output index arity.""" | |
| wf = hb.fill_workflow(scenes[0], artwork_dir) | |
| for node in wf.values(): | |
| for v in node["inputs"].values(): | |
| if isinstance(v, list) and len(v) == 2 and isinstance(v[0], str): | |
| assert v[0] in wf, f"dangling reference {v}" | |
| # KSampler consumes WanImageToVideo's 3 outputs (pos/neg/latent) | |
| assert wf["46"]["inputs"]["positive"] == ["45", 0] | |
| assert wf["46"]["inputs"]["negative"] == ["45", 1] | |
| assert wf["46"]["inputs"]["latent_image"] == ["45", 2] | |
| def test_estimate_cost_scales(): | |
| e1 = hb.estimate_cost([{}] * 10) | |
| e2 = hb.estimate_cost([{}] * 50) | |
| assert e2["estimated_cost_usd"] > e1["estimated_cost_usd"] | |
| assert e1["rate_per_hour"] == 1.28 | |
| def test_bootstrap_embeds_hf_repo(): | |
| assert "Comfy-Org/Wan_2.1_ComfyUI_repackaged" in hb.BOOTSTRAP | |
| assert "COMFY_REPO" not in hb.BOOTSTRAP | |
| assert "main.py --listen 127.0.0.1" in hb.BOOTSTRAP | |
| def test_constants_match_proven_recipe(): | |
| assert hb.BASE == "https://infrahub-api.nexgencloud.com/v1" | |
| assert hb.FLAVOR_SPOT == "n3-A100x1-spot" | |
| assert hb.ENV_NAME == "default-CANADA-1" | |
| def test_no_secrets_in_source(): | |
| """Guardrail: no hardcoded key value, only env/file sourcing of the key name.""" | |
| src = Path(hb.__file__).read_text() | |
| assert not re.search(r"HYPERSTACK_API_KEY\s*=\s*[\"']?[A-Za-z0-9_\-]{20,}", src) | |
| assert not re.search(r"[\"']sk-[A-Za-z0-9]{16,}[\"']", src) | |
| def test_build_manifest_from_dir(artwork_dir): | |
| m = hb.build_manifest_from_dir(artwork_dir) | |
| assert len(m) == 3 | |
| assert [s["artwork"] for s in m] == ["cover_1.png", "cover_2.png", "cover_3.png"] | |
| assert [s["scene_num"] for s in m] == [1, 2, 3] | |
| assert all(s["prompt"] for s in m) | |
| def test_build_manifest_custom_prompt_duration(artwork_dir): | |
| (artwork_dir / "sketch.txt").write_text("ignore me") | |
| m = hb.build_manifest_from_dir(artwork_dir, prompt="drift", duration=6, resolution="720p") | |
| assert len(m) == 3 # .txt filtered out | |
| assert m[0]["prompt"] == "drift" and m[0]["duration"] == 6 and m[0]["resolution"] == "720p" | |
| def test_build_manifest_empty_dir_raises(tmp_path): | |
| d = tmp_path / "empty" | |
| d.mkdir() | |
| with pytest.raises(ValueError): | |
| hb.build_manifest_from_dir(d) | |
| def test_notify_is_silent_no_crash(): | |
| hb.notify("t", "b") # osascript missing/blocked must not raise | |
| def test_run_burst_detached_spawns_child(tmp_path, monkeypatch): | |
| """Detach writes manifest + spawns the CLI; we stub Popen to avoid a real run.""" | |
| launched = {} | |
| class FakePopen: | |
| def __init__(self, cmd, **kw): | |
| launched["cmd"] = cmd | |
| monkeypatch.setattr(hb.subprocess, "Popen", FakePopen) | |
| scenes = [{"scene_num": 1, "artwork": "x.png", "prompt": "p"}] | |
| hb.run_burst_detached(scenes, tmp_path, tmp_path / "out") | |
| assert "--go" in launched["cmd"] and "--notify-done" in launched["cmd"] | |
| saved = json.loads((hb.STATE_DIR / "manifest.json").read_text()) | |
| assert saved == scenes | |
| def test_check_cost_guard_blocks_over_cap(): | |
| est = {"estimated_cost_usd": 4.0, "estimated_minutes": 225, "scenes": 50, "rate_per_hour": 1.28} | |
| with pytest.raises(SystemExit): | |
| hb.check_cost_guard(est, 3.0) | |
| def test_check_cost_guard_allows_under_cap(): | |
| est = {"estimated_cost_usd": 1.39, "estimated_minutes": 65, "scenes": 10, "rate_per_hour": 1.28} | |
| hb.check_cost_guard(est, 5.0) # no raise | |
| def test_model_sampling_sd3_matches_official_template(): | |
| """Official Comfy-Org Wan2.1 I2V template routes UNET through ModelSamplingSD3 (shift=8).""" | |
| wf = hb.fill_workflow({"scene_num": 1, "artwork": "cover_1.png", "prompt": "x"}, Path(".")) | |
| assert wf["36"]["class_type"] == "ModelSamplingSD3" | |
| assert wf["36"]["inputs"]["model"] == ["37", 0] | |
| assert wf["36"]["inputs"]["shift"] == 8 | |
| assert wf["46"]["inputs"]["model"] == ["36", 0] # KSampler consumes shifted model | |
| def test_official_template_crosscheck(): | |
| """Cross-check node types and key widgets against the official template file if cached.""" | |
| official = Path.home() / ".flomaster" / "scratch" / "wan_official.json" | |
| if not official.exists(): | |
| import pytest | |
| pytest.skip("official template not cached") | |
| d = json.loads(official.read_text())["nodes"] | |
| off_types = {n["type"] for n in d if n.get("type") != "MarkdownNote"} | |
| wf = hb.fill_workflow({"scene_num": 1, "artwork": "cover_1.png", "prompt": "x"}, Path(".")) | |
| my_types = {v["class_type"] for v in wf.values()} | |
| missing = off_types - my_types | |
| assert not missing, f"workflow missing official node types: {missing}" | |