| 118 | class DiTBlock(Module): |
| 119 | |
| 120 | def __init__(self, |
| 121 | hidden_size, |
| 122 | num_heads, |
| 123 | mapping=Mapping(), |
| 124 | mlp_ratio=4.0, |
| 125 | dtype=None, |
| 126 | quant_mode=QuantMode(0)): |
| 127 | super().__init__() |
| 128 | self.dtype = dtype |
| 129 | self.norm1 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 130 | self.attn = BertAttention(hidden_size, |
| 131 | num_heads, |
| 132 | tp_group=mapping.tp_group, |
| 133 | tp_size=mapping.tp_size, |
| 134 | tp_rank=mapping.tp_rank, |
| 135 | cp_group=mapping.cp_group, |
| 136 | cp_size=mapping.cp_size, |
| 137 | cp_rank=mapping.cp_rank, |
| 138 | dtype=dtype, |
| 139 | quant_mode=quant_mode) |
| 140 | self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 141 | self.mlp = MLP(hidden_size=hidden_size, |
| 142 | ffn_hidden_size=int(hidden_size * mlp_ratio), |
| 143 | hidden_act='gelu', |
| 144 | tp_group=mapping.tp_group, |
| 145 | tp_size=mapping.tp_size, |
| 146 | dtype=dtype, |
| 147 | quant_mode=quant_mode) |
| 148 | self.adaLN_modulation = Linear(hidden_size, |
| 149 | 6 * hidden_size, |
| 150 | tp_group=mapping.tp_group, |
| 151 | tp_size=mapping.tp_size, |
| 152 | bias=True, |
| 153 | dtype=dtype) |
| 154 | |
| 155 | def forward(self, x, c, input_lengths): |
| 156 | c = self.adaLN_modulation(silu(c)) |