| 124 | |
| 125 | |
| 126 | class AttentionBlock(torch.nn.Module): |
| 127 | |
| 128 | def __init__(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, cross_attention_dim=None, norm_num_groups=32, eps=1e-5, need_proj_out=True): |
| 129 | super().__init__() |
| 130 | inner_dim = num_attention_heads * attention_head_dim |
| 131 | |
| 132 | self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True) |
| 133 | self.proj_in = torch.nn.Linear(in_channels, inner_dim) |
| 134 | |
| 135 | self.transformer_blocks = torch.nn.ModuleList([ |
| 136 | BasicTransformerBlock( |
| 137 | inner_dim, |
| 138 | num_attention_heads, |
| 139 | attention_head_dim, |
| 140 | cross_attention_dim=cross_attention_dim |
| 141 | ) |
| 142 | for d in range(num_layers) |
| 143 | ]) |
| 144 | self.need_proj_out = need_proj_out |
| 145 | if need_proj_out: |
| 146 | self.proj_out = torch.nn.Linear(inner_dim, in_channels) |
| 147 | |
| 148 | def forward( |
| 149 | self, |
| 150 | hidden_states, time_emb, text_emb, res_stack, |
| 151 | cross_frame_attention=False, |
| 152 | tiled=False, tile_size=64, tile_stride=32, |
| 153 | ipadapter_kwargs_list={}, |
| 154 | **kwargs |
| 155 | ): |
| 156 | batch, _, height, width = hidden_states.shape |
| 157 | residual = hidden_states |
| 158 | |
| 159 | hidden_states = self.norm(hidden_states) |
| 160 | inner_dim = hidden_states.shape[1] |
| 161 | hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim) |
| 162 | hidden_states = self.proj_in(hidden_states) |
| 163 | |
| 164 | if cross_frame_attention: |
| 165 | hidden_states = hidden_states.reshape(1, batch * height * width, inner_dim) |
| 166 | encoder_hidden_states = text_emb.mean(dim=0, keepdim=True) |
| 167 | else: |
| 168 | encoder_hidden_states = text_emb |
| 169 | if encoder_hidden_states.shape[0] != hidden_states.shape[0]: |
| 170 | encoder_hidden_states = encoder_hidden_states.repeat(hidden_states.shape[0], 1, 1) |
| 171 | |
| 172 | if tiled: |
| 173 | tile_size = min(tile_size, min(height, width)) |
| 174 | hidden_states = hidden_states.permute(0, 2, 1).reshape(batch, inner_dim, height, width) |
| 175 | def block_tile_forward(x): |
| 176 | b, c, h, w = x.shape |
| 177 | x = x.permute(0, 2, 3, 1).reshape(b, h*w, c) |
| 178 | x = block(x, encoder_hidden_states) |
| 179 | x = x.reshape(b, h, w, c).permute(0, 3, 1, 2) |
| 180 | return x |
| 181 | for block in self.transformer_blocks: |
| 182 | hidden_states = TileWorker().tiled_forward( |
| 183 | block_tile_forward, |