(
self,
in_channels,
hidden_size,
heads_num,
depth,
mlp_width_ratio: float = 4.0,
mlp_drop_rate: float = 0.0,
act_type: str = "silu",
qk_norm: bool = False,
qk_norm_type: str = "layer",
qkv_bias: bool = True,
need_CA:bool=False,
attn_mode: str = "torch",
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
)
| 549 | A single token refiner block for llm text embedding refine. |
| 550 | """ |
| 551 | def __init__( |
| 552 | self, |
| 553 | in_channels, |
| 554 | hidden_size, |
| 555 | heads_num, |
| 556 | depth, |
| 557 | mlp_width_ratio: float = 4.0, |
| 558 | mlp_drop_rate: float = 0.0, |
| 559 | act_type: str = "silu", |
| 560 | qk_norm: bool = False, |
| 561 | qk_norm_type: str = "layer", |
| 562 | qkv_bias: bool = True, |
| 563 | need_CA:bool=False, |
| 564 | attn_mode: str = "torch", |
| 565 | dtype: Optional[torch.dtype] = None, |
| 566 | device: Optional[torch.device] = None, |
| 567 | ): |
| 568 | factory_kwargs = {"device": device, "dtype": dtype} |
| 569 | super().__init__() |
| 570 | self.attn_mode = attn_mode |
| 571 | self.need_CA = need_CA |
| 572 | assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner." |
| 573 | |
| 574 | self.input_embedder = nn.Linear( |
| 575 | in_channels, hidden_size, bias=True, **factory_kwargs |
| 576 | ) |
| 577 | if self.need_CA: |
| 578 | self.input_embedder_CA = nn.Linear( |
| 579 | in_channels, hidden_size, bias=True, **factory_kwargs |
| 580 | ) |
| 581 | |
| 582 | act_layer = get_activation_layer(act_type) |
| 583 | # Build timestep embedding layer |
| 584 | self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs) |
| 585 | # Build context embedding layer |
| 586 | self.c_embedder = TextProjection( |
| 587 | in_channels, hidden_size, act_layer, **factory_kwargs |
| 588 | ) |
| 589 | |
| 590 | self.individual_token_refiner = IndividualTokenRefiner( |
| 591 | hidden_size=hidden_size, |
| 592 | heads_num=heads_num, |
| 593 | depth=depth, |
| 594 | mlp_width_ratio=mlp_width_ratio, |
| 595 | mlp_drop_rate=mlp_drop_rate, |
| 596 | act_type=act_type, |
| 597 | qk_norm=qk_norm, |
| 598 | qk_norm_type=qk_norm_type, |
| 599 | qkv_bias=qkv_bias, |
| 600 | need_CA=need_CA, |
| 601 | **factory_kwargs, |
| 602 | ) |
| 603 | |
| 604 | def forward( |
| 605 | self, |
nothing calls this directly
no test coverage detected