File size: 5,942 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 | #!/usr/bin/env python3
"""CPU regression for the two-ended H3 AV bridge target and masks."""
import importlib.util
import os
import sys
import types
import torch
import torch.nn.functional as functional
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PACKAGE = "h3_bridge_test_pkg"
class NestedTensor:
def __init__(self, parts):
self.parts = tuple(parts)
def unbind(self):
return self.parts
def _install_stubs():
package = types.ModuleType(PACKAGE)
package.__path__ = [ROOT]
sys.modules[PACKAGE] = package
comfy = types.ModuleType("comfy")
comfy.__path__ = []
nested = types.ModuleType("comfy.nested_tensor")
nested.NestedTensor = NestedTensor
comfy.nested_tensor = nested
sys.modules["comfy"] = comfy
sys.modules["comfy.nested_tensor"] = nested
nodes = types.ModuleType("%s.nodes" % PACKAGE)
nodes.FPS = 24
nodes.AUDIO_HZ = 40
frame_per_token = (1, 4, 4, 4, 4)
nodes._pixel_frames = lambda count: sum(
frame_per_token[index % 5] for index in range(int(count)))
def resize(images, width, height, _crop):
work = images[..., :3].movedim(-1, 1)
work = functional.interpolate(
work, size=(height, width), mode="bilinear", align_corners=False)
return work.movedim(1, -1)
nodes._resize = resize
sys.modules[nodes.__name__] = nodes
support = types.ModuleType("%s.masking_support" % PACKAGE)
support.require_h3_mask_support = lambda _purpose: None
sys.modules[support.__name__] = support
return nodes
def _load_bridge():
spec = importlib.util.spec_from_file_location(
"%s.masked_bridge" % PACKAGE,
os.path.join(ROOT, "masked_bridge.py"))
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
class VideoVAE:
def encode(self, frames):
count = int(frames.shape[0])
steps = 2 if count <= 5 else ((count - 5) // 17) * 5 + 2
value = float(frames.mean())
return torch.full((1, 24, steps, 2, 4), value)
class AudioVAE:
audio_sample_rate = 32000
def encode(self, audio):
steps = round(int(audio.shape[1]) / self.audio_sample_rate * 40)
return torch.full((1, 32, 2, steps), float(audio.mean()))
def main():
nodes = _install_stubs()
bridge = _load_bridge()
assert nodes._pixel_frames(57) == 192
target_video = torch.zeros((1, 24, 57, 2, 4))
target_audio = torch.zeros((1, 32, 2, 320))
target = {"samples": NestedTensor((target_video, target_audio))}
start_frames = torch.full((100, 32, 64, 3), 0.2)
end_frames = torch.full((100, 32, 64, 3), 0.8)
start_audio = {
"waveform": torch.full(
(1, 2, round(100 / 24 * 32000)), 0.3),
"sample_rate": 32000,
}
end_audio = {
"waveform": torch.full(
(1, 2, round(100 / 24 * 32000)), 0.7),
"sample_rate": 32000,
}
input_types = bridge.MiniMaxH3ContexMaskedAVBridge.INPUT_TYPES()
assert "start_audio" not in input_types["required"]
assert "end_audio" not in input_types["required"]
assert {"start_audio", "end_audio"} <= set(input_types["optional"])
output, middle, protected = (
bridge.MiniMaxH3ContexMaskedAVBridge().prepare(
target,
VideoVAE(),
AudioVAE(),
start_frames,
start_audio,
end_frames,
end_audio,
24.0,
24.0,
39,
"disabled",
))
assert middle == 114
assert protected == 39
video, audio = output["samples"].unbind()
video_mask, audio_mask = output["noise_mask"].unbind()
assert torch.allclose(
video[:, :, :12], torch.full_like(video[:, :, :12], 0.2))
assert not torch.count_nonzero(video[:, :, 12:-12])
assert torch.allclose(
video[:, :, -12:], torch.full_like(video[:, :, -12:], 0.8))
assert torch.allclose(
audio[..., :65], torch.full_like(audio[..., :65], 0.3))
assert not torch.count_nonzero(audio[..., 65:-65])
assert torch.allclose(
audio[..., -65:], torch.full_like(audio[..., -65:], 0.7))
assert not torch.count_nonzero(video_mask[:, :, :12])
assert torch.all(video_mask[:, :, 12:-12] == 1)
assert not torch.count_nonzero(video_mask[:, :, -12:])
assert not torch.count_nonzero(audio_mask[..., :65])
assert torch.all(audio_mask[..., 65:-65] == 1)
assert not torch.count_nonzero(audio_mask[..., -65:])
assert not torch.count_nonzero(target_video)
assert not torch.count_nonzero(target_audio)
generated_start, _, _ = (
bridge.MiniMaxH3ContexMaskedAVBridge().prepare(
target, VideoVAE(), AudioVAE(), start_frames, None,
end_frames, end_audio, 24.0, 24.0, 39, "disabled"))
_generated_video, generated_audio = generated_start["samples"].unbind()
_generated_video_mask, generated_audio_mask = (
generated_start["noise_mask"].unbind())
assert not torch.count_nonzero(generated_audio[..., :65])
assert torch.all(generated_audio_mask[..., :65] == 1)
assert torch.allclose(
generated_audio[..., -65:],
torch.full_like(generated_audio[..., -65:], 0.7))
assert not torch.count_nonzero(generated_audio_mask[..., -65:])
try:
bridge.MiniMaxH3ContexMaskedAVBridge().prepare(
target, VideoVAE(), AudioVAE(), start_frames, start_audio,
end_frames, end_audio, 24.0, 24.0, 40, "disabled")
except ValueError as exc:
assert "exact H3 video run" in str(exc)
else:
raise AssertionError("a 40-frame bridge endpoint must be rejected")
print(
"masked bridge: 192 target frames, 39+39 protected, 114 generated; "
"12 video and optional 65-step audio windows protected independently")
if __name__ == "__main__":
main()
|