(
self,
in_channels,
n_heads,
d_head,
depth=1,
dropout=0.0,
use_linear=False,
context_dim=None,
use_spatial_context=False,
timesteps=None,
merge_strategy: str = "fixed",
merge_factor: float = 0.5,
time_context_dim=None,
ff_in=False,
checkpoint=False,
time_depth=1,
attn_mode="softmax",
disable_self_attn=False,
disable_temporal_crossattention=False,
max_time_embed_period: int = 10000,
dtype="fp32",
)
| 142 | |
| 143 | class SpatialVideoTransformer(SpatialTransformer): |
| 144 | def __init__( |
| 145 | self, |
| 146 | in_channels, |
| 147 | n_heads, |
| 148 | d_head, |
| 149 | depth=1, |
| 150 | dropout=0.0, |
| 151 | use_linear=False, |
| 152 | context_dim=None, |
| 153 | use_spatial_context=False, |
| 154 | timesteps=None, |
| 155 | merge_strategy: str = "fixed", |
| 156 | merge_factor: float = 0.5, |
| 157 | time_context_dim=None, |
| 158 | ff_in=False, |
| 159 | checkpoint=False, |
| 160 | time_depth=1, |
| 161 | attn_mode="softmax", |
| 162 | disable_self_attn=False, |
| 163 | disable_temporal_crossattention=False, |
| 164 | max_time_embed_period: int = 10000, |
| 165 | dtype="fp32", |
| 166 | ): |
| 167 | super().__init__( |
| 168 | in_channels, |
| 169 | n_heads, |
| 170 | d_head, |
| 171 | depth=depth, |
| 172 | dropout=dropout, |
| 173 | attn_type=attn_mode, |
| 174 | use_checkpoint=checkpoint, |
| 175 | context_dim=context_dim, |
| 176 | use_linear=use_linear, |
| 177 | disable_self_attn=disable_self_attn, |
| 178 | ) |
| 179 | self.time_depth = time_depth |
| 180 | self.depth = depth |
| 181 | self.max_time_embed_period = max_time_embed_period |
| 182 | |
| 183 | time_mix_d_head = d_head |
| 184 | n_time_mix_heads = n_heads |
| 185 | |
| 186 | time_mix_inner_dim = int(time_mix_d_head * n_time_mix_heads) |
| 187 | |
| 188 | inner_dim = n_heads * d_head |
| 189 | if use_spatial_context: |
| 190 | time_context_dim = context_dim |
| 191 | |
| 192 | self.time_stack = nn.ModuleList( |
| 193 | [ |
| 194 | VideoTransformerBlock( |
| 195 | inner_dim, |
| 196 | n_time_mix_heads, |
| 197 | time_mix_d_head, |
| 198 | dropout=dropout, |
| 199 | context_dim=time_context_dim, |
| 200 | timesteps=timesteps, |
| 201 | checkpoint=checkpoint, |
nothing calls this directly
no test coverage detected