Spaces:
Running on Zero
Running on Zero
File size: 5,574 Bytes
49bc52e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | from typing import Tuple
from einops import rearrange
import torch
from torch import nn
from .wan_video_dit import GateModule
def create_patch_embedding(in_channels, out_channels, patch_kernel_size, wan_patch_embedding=None):
patch_embedding = nn.Conv3d(in_channels, out_channels, kernel_size=patch_kernel_size, stride=patch_kernel_size)
if wan_patch_embedding is not None:
target_out_channels, target_in_channels, kernel_1, kernel_2, kernel_3 = wan_patch_embedding.weight.shape
target_patch_kernel_size = (kernel_1, kernel_2, kernel_3)
assert out_channels == target_out_channels,\
f"`out_channels` ({out_channels}) must match that of `wan_patch_embedding` ({target_out_channels})"
assert in_channels >= target_in_channels,\
f"`in_channels` ({in_channels}) must be at least that of `wan_patch_embedding` ({target_in_channels})"
assert patch_kernel_size == target_patch_kernel_size, (
f"`patch_kernel_size` ({patch_kernel_size}) must match that of "
f"`wan_patch_embedding` ({target_patch_kernel_size})"
)
dtype = wan_patch_embedding.weight.dtype
device = wan_patch_embedding.weight.device
weight_additional = torch.zeros(
out_channels, in_channels - target_in_channels, *target_patch_kernel_size, dtype=dtype, device=device,
)
weight = torch.cat([wan_patch_embedding.weight.data.clone(), weight_additional], dim=1)
assert weight.shape == patch_embedding.weight.shape,\
f"Shape of new weight {weight.shape} must match that of original weight {patch_embedding.weight.shape}"
patch_embedding.weight.data = nn.Parameter(weight)
patch_embedding.bias.data = nn.Parameter(wan_patch_embedding.bias.data.clone())
else:
nn.init.zeros_(patch_embedding.weight)
nn.init.zeros_(patch_embedding.bias)
return patch_embedding
def patchify(x, patch_embedding):
x = patch_embedding(x)
grid_size = x.shape[2:]
x = rearrange(x, "b c f h w -> b (f h w) c").contiguous()
return x, grid_size # x, grid_size: (f, h, w)
class LatentEncoder(nn.Module):
def __init__(
self,
in_channels: int = 16,
num_inputs: int = 2,
hidden_channels: int = 1024,
out_channels: int = 5120,
patch_kernel_size: Tuple[int] = (1, 2, 2),
wan_patch_embedding: nn.Module = None,
target_reuse_patch_embedding: bool = False,
):
super().__init__()
assert wan_patch_embedding is not None
self.patch_embedding = lambda x: wan_patch_embedding(x)
self.target_reuse_patch_embedding = target_reuse_patch_embedding
dtype = wan_patch_embedding.weight.dtype
device = wan_patch_embedding.weight.device
# Reference (masked source trajectory)
self.reference_patch_embedding = create_patch_embedding(
in_channels * num_inputs, out_channels, patch_kernel_size, wan_patch_embedding=wan_patch_embedding,
)
# Target (masked target trajectory)
if target_reuse_patch_embedding:
self.target_patch_embedding = create_patch_embedding(
in_channels * num_inputs, out_channels, patch_kernel_size,
wan_patch_embedding=wan_patch_embedding,
)
self.gate = GateModule()
self.target_gate = nn.Parameter(torch.zeros(out_channels, 1, 1, 1, dtype=dtype, device=device))
else:
self.target_encoder = nn.Sequential( # From EX-4D
nn.Linear(in_channels * num_inputs, hidden_channels, bias=True),
nn.SiLU(),
nn.Linear(hidden_channels, hidden_channels, bias=True),
nn.SiLU(),
nn.Linear(hidden_channels, hidden_channels, bias=True),
)
self.target_patch_embedding = nn.Conv3d(
hidden_channels, out_channels, kernel_size=patch_kernel_size, stride=patch_kernel_size,
)
nn.init.zeros_(self.target_patch_embedding.weight)
nn.init.zeros_(self.target_patch_embedding.bias)
def forward(self, x, reference_video_latents, reference_mask_latents, target_video_latents, target_mask_latents):
x, (f, h, w) = patchify(x, self.patch_embedding)
masked_reference = None
if reference_video_latents is not None and reference_mask_latents is not None:
masked_reference = torch.cat([reference_video_latents, reference_mask_latents], dim=-4)
masked_reference, (f_, h_, w_) = patchify(masked_reference, self.reference_patch_embedding)
assert (f, h, w) == (f_, h_, w_)
if target_video_latents is not None and target_mask_latents is not None:
masked_target = torch.cat([target_video_latents, target_mask_latents], dim=-4)
if self.target_reuse_patch_embedding:
masked_target, _ = patchify(masked_target, self.target_patch_embedding)
x = self.target_gate_apply(x, self.target_gate, masked_target)
else:
masked_target = masked_target.moveaxis(-4, -1) # b c f h w -> b f h w c
masked_target = self.target_encoder(masked_target)
masked_target = masked_target.moveaxis(-1, -4) # b f h w c -> b c f h w
masked_target, (f_, h_, w_) = patchify(masked_target, self.target_patch_embedding)
assert (f, h, w) == (f_, h_, w_)
x = x + masked_target
return x, masked_reference, (f, h, w)
|