File size: 10,183 Bytes
bc50ea0 | 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 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 | #!/usr/bin/env python3
"""Chapter delivery, selection, isolation, and recovery regressions."""
import hashlib
import importlib.util
import pathlib
import re
import sys
import tempfile
import types
import torch
ROOT = pathlib.Path(__file__).resolve().parents[1]
PACKAGE = "h3_chapter_delivery_unit"
folder_paths = types.ModuleType("folder_paths")
folder_paths.output_directory = str(ROOT)
folder_paths.get_output_directory = lambda: folder_paths.output_directory
folder_paths.get_temp_directory = lambda: folder_paths.output_directory
folder_paths.get_input_directory = lambda: folder_paths.output_directory
folder_paths.get_annotated_filepath = lambda value: str(value)
sys.modules["folder_paths"] = folder_paths
package = types.ModuleType(PACKAGE)
package.__path__ = [str(ROOT)]
sys.modules[PACKAGE] = package
shared_nodes = types.ModuleType(PACKAGE + ".nodes")
shared_nodes.MiniMaxH3MotionContext = object
shared_nodes._claim_inline_patch_ownership = lambda _conditioning=None: "test"
shared_nodes._prepare_native_guide_conditioning = lambda value: value
shared_nodes._resize = lambda *args: None
shared_nodes._streams_from_latent = lambda *args: None
sys.modules[shared_nodes.__name__] = shared_nodes
spec = importlib.util.spec_from_file_location(
PACKAGE + ".chain_nodes", ROOT / "chain_nodes.py")
chain = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = chain
spec.loader.exec_module(chain)
def sha256(path):
return hashlib.sha256(path.read_bytes()).hexdigest()
def make_segment(output, index, delivered=10):
run = output / "h3_chains" / "chapter_delivery"
segment_path = run / "segments" / ("clip_%04d.mp4" % index)
checkpoint_path = run / "checkpoints" / ("clip_%04d.safetensors" % index)
segment_path.parent.mkdir(parents=True, exist_ok=True)
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
segment_path.write_bytes(("video-%d" % index).encode("ascii"))
checkpoint_path.write_bytes(("latent-%d" % index).encode("ascii"))
return {
"index": index,
"id": "scene_%02d" % index,
"revision": ("%032x" % index),
"segment": str(segment_path.relative_to(output)),
"segment_sha256": sha256(segment_path),
"checkpoint": str(checkpoint_path.relative_to(output)),
"checkpoint_sha256": sha256(checkpoint_path),
"raw_frames": delivered,
"delivered_frames": delivered,
"blend_frames": 0,
}
def make_manifest(segments, complete):
total = sum(int(item["delivered_frames"]) for item in segments)
result = {
"format": ("h3_chain_manifest_v3" if complete else
"h3_chain_partial_manifest_v3"),
"run_name": "chapter_delivery",
"plan_hash": "plan-hash",
"prompt_prefix": "test",
"compatibility": {"continuation_mode": "guide"},
"clip_count": len(segments),
"total_delivered_frames": total,
"duration_seconds": total / float(chain.FPS),
"segments": segments,
}
if not complete:
result.update({
"planned_clip_count": 6,
"last_completed_clip": len(segments),
})
return result
def make_editorial():
order = [{"scene": index, "scene_id": "scene_%02d" % index}
for index in range(1, 7)]
return {
"format": "h3_chain_editorial_v1",
"run_name": "chapter_delivery",
"scene_order": order,
"chapters": [
{"id": "opening", "title": "Opening", "start_scene": 1,
"start_scene_id": "scene_01", "text": "First delivery"},
{"id": "aftermath", "title": "Aftermath", "start_scene": 4,
"start_scene_id": "scene_04", "text": "Second delivery"},
],
"placements": [],
"trims": [],
"locked_scene_ids": [],
"subtitles": {"mode": "off", "asset_id": "",
"offset_seconds": 0.0},
"alternate_draft": None,
"replacements": [],
}
def expect_error(action, phrase):
try:
action()
except (FileNotFoundError, ValueError) as exc:
assert phrase in str(exc), str(exc)
else:
raise AssertionError("expected error containing %r" % phrase)
def main():
assert chain.CHAIN_NODE_CLASS_MAPPINGS[
"MiniMaxH3ChainChapterDelivery"] is chain.MiniMaxH3ChainChapterDelivery
assert chain.CHAIN_NODE_CLASS_MAPPINGS[
"MiniMaxH3ChainChapterLoad"] is chain.MiniMaxH3ChainChapterLoad
schema = chain.MiniMaxH3ChainChapterDelivery.INPUT_TYPES()
assert schema["required"]["chapter_number"][1]["default"] == 0
with tempfile.TemporaryDirectory() as temporary:
output = pathlib.Path(temporary)
folder_paths.output_directory = temporary
segments = [make_segment(output, index) for index in range(1, 7)]
chain._atomic_json(
chain._run_editorial_path("chapter_delivery"), make_editorial())
partial_three = make_manifest(segments[:3], complete=False)
chapter_one, chapter_one_path = chain._chapter_manifest_from_manifest(
partial_three, 0)
assert chapter_one["format"] == chain.CHAPTER_MANIFEST_FORMAT
assert chapter_one["chapter"]["number"] == 1
assert chapter_one["scene_start"] == 1
assert chapter_one["scene_end"] == 3
assert chapter_one["chapter"]["complete"] is True
assert chapter_one["chapter"]["planned_end_scene"] == 3
assert [item["index"] for item in chapter_one["segments"]] == [1, 2, 3]
assert pathlib.Path(chapter_one_path).is_file()
assert "/chapters/01_opening/manifests/" in chapter_one_path
assert re.fullmatch(
r"[0-9a-f]{32}", chapter_one["chapter_manifest_id"])
chapter_one_again, chapter_one_again_path = (
chain._chapter_manifest_from_manifest(partial_three, 0))
assert (chapter_one_again["chapter_manifest_id"] ==
chapter_one["chapter_manifest_id"])
assert chapter_one_again_path == chapter_one_path
complete = make_manifest(segments, complete=True)
completed_chapter_one, _completed_chapter_one_path = (
chain._chapter_manifest_from_manifest(complete, 1))
assert completed_chapter_one["chapter"]["number"] == 1
assert [item["index"] for item in
completed_chapter_one["segments"]] == [1, 2, 3]
chapter_two, chapter_two_path = chain._chapter_manifest_from_manifest(
complete, 2)
assert chapter_two["chapter"]["number"] == 2
assert chapter_two["chapter"]["source_start_frame"] == 30
assert chapter_two["chapter"]["editorial_origin_frame"] == 30
assert chapter_two["scene_start"] == 4
assert chapter_two["scene_end"] == 6
assert [item["index"] for item in chapter_two["segments"]] == [4, 5, 6]
assert "prelude" not in chapter_two
assert "/chapters/02_aftermath/manifests/" in chapter_two_path
assert chain._validate_manifest(chapter_two) == chapter_two["segments"]
auto_two, _auto_two_path = chain._chapter_manifest_from_manifest(
complete, 0)
assert auto_two["chapter"]["number"] == 2
partial_four = make_manifest(segments[:4], complete=False)
partial_two, partial_two_path = chain._chapter_manifest_from_manifest(
partial_four, 0)
assert partial_two["chapter"]["number"] == 2
assert partial_two["chapter"]["complete"] is False
assert partial_two["chapter"]["planned_end_scene"] == 6
assert partial_two["scene_start"] == partial_two["scene_end"] == 4
assert chain._validate_manifest(partial_two) == partial_two["segments"]
explicit_one, _explicit_one_path = (
chain._chapter_manifest_from_manifest(partial_four, 1))
assert explicit_one["chapter"]["number"] == 1
explicit_two, explicit_two_path = chain._chapter_manifest_from_manifest(
partial_four, 2)
assert explicit_two == partial_two and explicit_two_path == partial_two_path
expect_error(
lambda: chain._chapter_manifest_from_manifest(partial_three, 2),
"no generated scenes yet")
expect_error(
lambda: chain._chapter_manifest_from_manifest(complete, 3),
"does not exist")
loaded, loaded_path = chain._load_chapter_manifest(
"chapter_delivery", 2)
assert loaded["chapter_manifest_id"] == partial_two[
"chapter_manifest_id"]
assert loaded_path == partial_two_path
exact, exact_path = chain._load_chapter_manifest(
"chapter_delivery", 2, chapter_two["chapter_manifest_id"])
assert exact == chapter_two and exact_path == chapter_two_path
media_metadata = chain._manifest_media_metadata(chapter_two)
assert "Chapter 2 Aftermath" in media_metadata["title"]
assert "Chapter 2 scenes 4-6" in media_metadata["comment"]
delivery = chain.MiniMaxH3ChainChapterDelivery()
passed, _json, number, path, status = delivery.select(
complete, enabled=False, chapter_number=2)
assert passed == complete and number == 0 and path == ""
assert "disabled" in status
export_dir = chain._new_export_directory(chapter_two, "plates")
assert "/chapters/02_aftermath/frames/plates" in export_dir
source_audio = {
"waveform": torch.arange(60, dtype=torch.float32).reshape(1, 1, 60),
"sample_rate": 24,
}
chapter_two["compatibility"]["source_audio_hash"] = (
chain._audio_fingerprint(source_audio))
selected = chain._full_chain_selected_audio(
chapter_two, "source", None, source_audio)
assert torch.equal(
selected["waveform"], source_audio["waveform"][..., 30:60])
expect_error(
lambda: chain._load_chapter_manifest("chapter_delivery", 3),
"No sealed Chapter 3")
print("H3 chapter delivery: explicit/auto selection, isolated output, "
"source-audio offset, immutable recovery, and pass-through passed")
if __name__ == "__main__":
main()
|