hugging-apps/echo-memory

1import torch2import torch.nn as nn3from .wan_video_dit import sinusoidal_embedding_1d4 5 6 7class WanMotionControllerModel(torch.nn.Module):8    def __init__(self, freq_dim=256, dim=1536):9        super().__init__()10        self.freq_dim = freq_dim11        self.linear = nn.Sequential(12            nn.Linear(freq_dim, dim),13            nn.SiLU(),14            nn.Linear(dim, dim),15            nn.SiLU(),16            nn.Linear(dim, dim * 6),17        )18 19    def forward(self, motion_bucket_id):20        emb = sinusoidal_embedding_1d(self.freq_dim, motion_bucket_id * 10)21        emb = self.linear(emb)22        return emb23 24    def init(self):25        state_dict = self.linear[-1].state_dict()26        state_dict = {i: state_dict[i] * 0 for i in state_dict}27        self.linear[-1].load_state_dict(state_dict)28 29    @staticmethod30    def state_dict_converter():31        return WanMotionControllerModelDictConverter()32    33    34 35class WanMotionControllerModelDictConverter:36    def __init__(self):37        pass38 39    def from_diffusers(self, state_dict):40        return state_dict41    42    def from_civitai(self, state_dict):43        return state_dict44 45