ALSv/self-forcing

1 2import torch3 4 5ZERO_VAE_CACHE = [6    torch.zeros(1, 16, 2, 60, 104),7    torch.zeros(1, 384, 2, 60, 104),8    torch.zeros(1, 384, 2, 60, 104),9    torch.zeros(1, 384, 2, 60, 104),10    torch.zeros(1, 384, 2, 60, 104),11    torch.zeros(1, 384, 2, 60, 104),12    torch.zeros(1, 384, 2, 60, 104),13    torch.zeros(1, 384, 2, 60, 104),14    torch.zeros(1, 384, 2, 60, 104),15    torch.zeros(1, 384, 2, 60, 104),16    torch.zeros(1, 384, 2, 60, 104),17    torch.zeros(1, 384, 2, 60, 104),18    torch.zeros(1, 192, 2, 120, 208),19    torch.zeros(1, 384, 2, 120, 208),20    torch.zeros(1, 384, 2, 120, 208),21    torch.zeros(1, 384, 2, 120, 208),22    torch.zeros(1, 384, 2, 120, 208),23    torch.zeros(1, 384, 2, 120, 208),24    torch.zeros(1, 384, 2, 120, 208),25    torch.zeros(1, 192, 2, 240, 416),26    torch.zeros(1, 192, 2, 240, 416),27    torch.zeros(1, 192, 2, 240, 416),28    torch.zeros(1, 192, 2, 240, 416),29    torch.zeros(1, 192, 2, 240, 416),30    torch.zeros(1, 192, 2, 240, 416),31    torch.zeros(1, 96, 2, 480, 832),32    torch.zeros(1, 96, 2, 480, 832),33    torch.zeros(1, 96, 2, 480, 832),34    torch.zeros(1, 96, 2, 480, 832),35    torch.zeros(1, 96, 2, 480, 832),36    torch.zeros(1, 96, 2, 480, 832),37    torch.zeros(1, 96, 2, 480, 832)38]39 40feat_names = [f"vae_cache_{i}" for i in range(len(ZERO_VAE_CACHE))]41ALL_INPUTS_NAMES = ["z", "use_cache"] + feat_names42