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