| |
| """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() |
|
|
