(
self,
*,
device: torch.device,
dtype: torch.dtype,
input_channels: int = 3,
output_channels: int = 3,
n_ctx: int = 1024,
width: int = 512,
layers: int = 12,
heads: int = 8,
init_scale: float = 0.25,
time_token_cond: bool = False,
)
| 154 | |
| 155 | class PointDiffusionTransformer(nn.Module): |
| 156 | def __init__( |
| 157 | self, |
| 158 | *, |
| 159 | device: torch.device, |
| 160 | dtype: torch.dtype, |
| 161 | input_channels: int = 3, |
| 162 | output_channels: int = 3, |
| 163 | n_ctx: int = 1024, |
| 164 | width: int = 512, |
| 165 | layers: int = 12, |
| 166 | heads: int = 8, |
| 167 | init_scale: float = 0.25, |
| 168 | time_token_cond: bool = False, |
| 169 | ): |
| 170 | super().__init__() |
| 171 | self.input_channels = input_channels |
| 172 | self.output_channels = output_channels |
| 173 | self.n_ctx = n_ctx |
| 174 | self.time_token_cond = time_token_cond |
| 175 | self.time_embed = MLP( |
| 176 | device=device, dtype=dtype, width=width, init_scale=init_scale * math.sqrt(1.0 / width) |
| 177 | ) |
| 178 | self.ln_pre = nn.LayerNorm(width, device=device, dtype=dtype) |
| 179 | self.backbone = Transformer( |
| 180 | device=device, |
| 181 | dtype=dtype, |
| 182 | n_ctx=n_ctx + int(time_token_cond), |
| 183 | width=width, |
| 184 | layers=layers, |
| 185 | heads=heads, |
| 186 | init_scale=init_scale, |
| 187 | ) |
| 188 | self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype) |
| 189 | self.input_proj = nn.Linear(input_channels, width, device=device, dtype=dtype) |
| 190 | self.output_proj = nn.Linear(width, output_channels, device=device, dtype=dtype) |
| 191 | with torch.no_grad(): |
| 192 | self.output_proj.weight.zero_() |
| 193 | self.output_proj.bias.zero_() |
| 194 | |
| 195 | def forward(self, x: torch.Tensor, t: torch.Tensor): |
| 196 | """ |
nothing calls this directly
no test coverage detected