(
self,
dim: int,
num_attention_heads: int,
attention_head_dim: int,
context_pre_only: bool = False,
qk_norm: str | None = None,
use_dual_attention: bool = False,
)
| 592 | """ |
| 593 | |
| 594 | def __init__( |
| 595 | self, |
| 596 | dim: int, |
| 597 | num_attention_heads: int, |
| 598 | attention_head_dim: int, |
| 599 | context_pre_only: bool = False, |
| 600 | qk_norm: str | None = None, |
| 601 | use_dual_attention: bool = False, |
| 602 | ): |
| 603 | super().__init__() |
| 604 | |
| 605 | self.use_dual_attention = use_dual_attention |
| 606 | self.context_pre_only = context_pre_only |
| 607 | context_norm_type = "ada_norm_continous" if context_pre_only else "ada_norm_zero" |
| 608 | |
| 609 | if use_dual_attention: |
| 610 | self.norm1 = SD35AdaLayerNormZeroX(dim) |
| 611 | else: |
| 612 | self.norm1 = AdaLayerNormZero(dim) |
| 613 | |
| 614 | if context_norm_type == "ada_norm_continous": |
| 615 | self.norm1_context = AdaLayerNormContinuous( |
| 616 | dim, dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type="layer_norm" |
| 617 | ) |
| 618 | elif context_norm_type == "ada_norm_zero": |
| 619 | self.norm1_context = AdaLayerNormZero(dim) |
| 620 | else: |
| 621 | raise ValueError( |
| 622 | f"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`" |
| 623 | ) |
| 624 | |
| 625 | if hasattr(F, "scaled_dot_product_attention"): |
| 626 | processor = JointAttnProcessor2_0() |
| 627 | else: |
| 628 | raise ValueError( |
| 629 | "The current PyTorch version does not support the `scaled_dot_product_attention` function." |
| 630 | ) |
| 631 | |
| 632 | self.attn = Attention( |
| 633 | query_dim=dim, |
| 634 | cross_attention_dim=None, |
| 635 | added_kv_proj_dim=dim, |
| 636 | dim_head=attention_head_dim, |
| 637 | heads=num_attention_heads, |
| 638 | out_dim=dim, |
| 639 | context_pre_only=context_pre_only, |
| 640 | bias=True, |
| 641 | processor=processor, |
| 642 | qk_norm=qk_norm, |
| 643 | eps=1e-6, |
| 644 | ) |
| 645 | |
| 646 | if use_dual_attention: |
| 647 | self.attn2 = Attention( |
| 648 | query_dim=dim, |
| 649 | cross_attention_dim=None, |
| 650 | dim_head=attention_head_dim, |
| 651 | heads=num_attention_heads, |
nothing calls this directly
no test coverage detected