()
| 80 | |
| 81 | |
| 82 | def load_custom_video_vae(): |
| 83 | with accelerate.init_empty_weights(): |
| 84 | vae = CausalContinuousVideoTokenizer( |
| 85 | attn_resolutions=[32], |
| 86 | channels=128, |
| 87 | channels_mult=[2, 4, 4], |
| 88 | dropout=0.0, |
| 89 | in_channels=3, |
| 90 | num_res_blocks=2, |
| 91 | out_channels=3, |
| 92 | resolution=1024, |
| 93 | patch_size=4, |
| 94 | patch_method="haar", |
| 95 | z_channels=16, |
| 96 | z_factor=1, |
| 97 | num_groups=1, |
| 98 | legacy_mode=False, |
| 99 | spatial_compression=8, |
| 100 | temporal_compression=8, |
| 101 | embedding_dim=16, |
| 102 | ) |
| 103 | missing_keys, unexpected_keys = vae.load_state_dict(load_state_dict(COMFYUI_VIDEO_VAE_WEIGHTS), assign=True, strict=False) |
| 104 | assert len(missing_keys) == 0 |
| 105 | vae.eval() |
| 106 | return vae |
| 107 | |
| 108 | |
| 109 | if __name__ == '__main__': |
no test coverage detected