(in_channels, attn_type="vanilla", attn_kwargs=None)
| 285 | |
| 286 | |
| 287 | def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None): |
| 288 | assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown' |
| 289 | if XFORMERS_IS_AVAILBLE and attn_type == "vanilla": |
| 290 | attn_type = "vanilla-xformers" |
| 291 | print(f"making attention of type '{attn_type}' with {in_channels} in_channels") |
| 292 | if attn_type == "vanilla": |
| 293 | assert attn_kwargs is None |
| 294 | return AttnBlock(in_channels) |
| 295 | elif attn_type == "vanilla-xformers": |
| 296 | print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...") |
| 297 | return MemoryEfficientAttnBlock(in_channels) |
| 298 | elif type == "memory-efficient-cross-attn": |
| 299 | attn_kwargs["query_dim"] = in_channels |
| 300 | return MemoryEfficientCrossAttentionWrapper(**attn_kwargs) |
| 301 | elif attn_type == "none": |
| 302 | return nn.Identity(in_channels) |
| 303 | else: |
| 304 | raise NotImplementedError() |
| 305 | |
| 306 | |
| 307 | class Model(nn.Module): |
no test coverage detected