Buckets:
| from unittest.mock import patch | |
| import pytest | |
| import torch | |
| import torch.nn as nn | |
| from comfy.cli_args import args as cli_args | |
| if not torch.cuda.is_available(): | |
| cli_args.cpu = True | |
| import comfy.ldm.seedvr.vae as vae_mod # noqa: E402 | |
| from comfy_extras import nodes_seedvr # noqa: E402 | |
| _LATENT_CHANNELS = vae_mod.SEEDVR2_LATENT_CHANNELS | |
| def _make_wrapper() -> vae_mod.VideoAutoencoderKLWrapper: | |
| wrapper = vae_mod.VideoAutoencoderKLWrapper.__new__( | |
| vae_mod.VideoAutoencoderKLWrapper | |
| ) | |
| nn.Module.__init__(wrapper) | |
| return wrapper | |
| def _fingerprint_decode(self, z, output_buffer=None): | |
| b = int(z.shape[0]) | |
| t = int(z.shape[2]) | |
| h = int(z.shape[3]) | |
| w = int(z.shape[4]) | |
| out = torch.empty(b, 3, t, h * 8, w * 8) | |
| for batch_idx in range(b): | |
| out[batch_idx].fill_(float(batch_idx + 1)) | |
| return out | |
| def _decode_with_patches(wrapper, z): | |
| with patch.object(vae_mod.VideoAutoencoderKLWrapper, "slicing_decode", _fingerprint_decode): | |
| return wrapper.decode(z) | |
| def test_decode_b2_t3_multi_frame_batch_unchanged(): | |
| wrapper = _make_wrapper() | |
| out = _decode_with_patches(wrapper, torch.zeros(2, _LATENT_CHANNELS * 3, 2, 2)) | |
| assert tuple(out.shape) == (2, 3, 3, 16, 16) | |
| class _Wrapper(vae_mod.VideoAutoencoderKLWrapper): | |
| def __init__(self): | |
| nn.Module.__init__(self) | |
| self.calls = [] | |
| def _decode_stub(self, latent, output_buffer=None): | |
| self.calls.append(tuple(latent.shape)) | |
| return torch.zeros(latent.shape[0], 3, latent.shape[2], latent.shape[3] * 8, latent.shape[4] * 8) | |
| def test_seedvr2_wrapper_decode_accepts_5d_channel_first_latents_without_preprocessor_state(): | |
| wrapper = _Wrapper() | |
| with patch.object(vae_mod.VideoAutoencoderKLWrapper, "slicing_decode", _decode_stub): | |
| out = wrapper.decode(torch.zeros(1, _LATENT_CHANNELS, 2, 4, 5)) | |
| assert tuple(out.shape) == (1, 3, 2, 32, 40) | |
| assert wrapper.calls == [(1, _LATENT_CHANNELS, 2, 4, 5)] | |
| def _t_padded(t_in: int) -> int: | |
| if t_in == 1: | |
| return 1 | |
| if t_in <= 4: | |
| return 5 | |
| if (t_in - 1) % 4 == 0: | |
| return t_in | |
| return t_in + (4 - ((t_in - 1) % 4)) | |
| def test_t_padded_matches_cut_videos(t_in): | |
| dummy = torch.zeros(1, t_in, 1, 1, 1) | |
| assert nodes_seedvr.cut_videos(dummy).shape[1] == _t_padded(t_in) | |
Xet Storage Details
- Size:
- 2.37 kB
- Xet hash:
- 83735abc80f9063f45bf01b2d922c804e5a524f2577b8566d8e71c5c6be915da
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.