(self, size, newtoken=1)
| 2 | import torch |
| 3 | class IterativeEmbedding(torch.nn.Module): |
| 4 | def __init__(self, size, newtoken=1): |
| 5 | super().__init__() |
| 6 | |
| 7 | timestep_input_dim = 768 |
| 8 | time_embed_dim = 2048 |
| 9 | flip_sin_to_cos = True |
| 10 | freq_shift = 0 |
| 11 | self.time_proj = Timesteps(timestep_input_dim, flip_sin_to_cos, freq_shift) |
| 12 | # timestep_input_dim = block_out_channels[0] |
| 13 | self.time_embedding = TimestepEmbedding( |
| 14 | timestep_input_dim, |
| 15 | time_embed_dim, |
| 16 | act_fn="silu", |
| 17 | out_dim=timestep_input_dim, |
| 18 | post_act_fn=None, |
| 19 | cond_proj_dim=None, |
| 20 | ) |
| 21 | # self.time_emb_proj = torch.nn.Linear(time_embed_dim, timestep_input_dim) |
| 22 | |
| 23 | # embed_dim = 768 |
| 24 | |
| 25 | # self.iterative_embedding = torch.nn.Parameter(torch.randn(768), requires_grad=True) |
| 26 | self.expand_embeddings = torch.nn.Embedding(size, 768) |
| 27 | self.expand_embeddings.weight.data.zero_() |
| 28 | self.expand_embeddings.weight.data[-newtoken:] = torch.ones(768) |
| 29 | self.expand_embeddings.weight.requires_grad = False |
| 30 | |
| 31 | |
| 32 | def forward( |
nothing calls this directly
no outgoing calls
no test coverage detected