Coverage for wrapper/flux2/transformer_flux2.py: 100%

1# Source (vendored, adapted): 

2# https://github.com/xdit-project/xDiT/blob/8f3e28a4f0c94545a1c2b9dfc3bb18b9e6f4c2d1/xfuser/model_executor/models/transformers/transformer_flux2.py 

3# See the upstream repository and its LICENSE file for original license and copyright details. 

4import torch 

5from typing import Any, Optional, Tuple 

6from diffusers.models.transformers.transformer_flux2 import ( 

7 Flux2Attention, 

8 Flux2AttnProcessor, 

9 Flux2Transformer2DModel, 

10 Flux2ParallelSelfAttention, 

11 Flux2ParallelSelfAttnProcessor, 

12 _get_qkv_projections, 

13) 

14from diffusers.models.embeddings import apply_rotary_emb 

15 

16from xfuser.model_executor.layers.attention_processor import ( 

17 xFuserAttentionBaseWrapper, 

18 xFuserAttentionProcessorRegister 

19) 

20 

21from xfuser.core.distributed import ( 

22 get_sequence_parallel_world_size, 

23 get_sequence_parallel_rank, 

24 get_classifier_free_guidance_world_size, 

25 get_classifier_free_guidance_rank, 

26 get_sp_group, 

27 get_cfg_group, 

28) 

29 

30from xfuser.model_executor.layers.usp import USP 

31from xfuser.model_executor.layers import xFuserLayerWrappersRegister 

32 

33 

34@xFuserAttentionProcessorRegister.register(Flux2AttnProcessor) 

35class xFuserFlux2AttnProcessor(Flux2AttnProcessor): 

36 

37 def __init__(self) -> None: 

38 super().__init__() 

39 

40 def __call__( 

41 self, 

42 attn: "Flux2Attention", 

43 hidden_states: torch.Tensor, 

44 encoder_hidden_states: Optional[torch.Tensor] = None, 

45 attention_mask: Optional[torch.Tensor] = None, 

46 image_rotary_emb: Optional[torch.Tensor] = None, 

47 ) -> Any: 

48 query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( 

49 attn, hidden_states, encoder_hidden_states 

50 ) 

51 

52 query = query.unflatten(-1, (attn.heads, -1)) 

53 key = key.unflatten(-1, (attn.heads, -1)) 

54 value = value.unflatten(-1, (attn.heads, -1)) 

55 

56 query = attn.norm_q(query) 

57 key = attn.norm_k(key) 

58 

59 if attn.added_kv_proj_dim is not None: 

60 encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) 

61 encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) 

62 encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) 

63 

64 encoder_query = attn.norm_added_q(encoder_query) 

65 encoder_key = attn.norm_added_k(encoder_key) 

66 

67 query = torch.cat([encoder_query, query], dim=1) 

68 key = torch.cat([encoder_key, key], dim=1) 

69 value = torch.cat([encoder_value, value], dim=1) 

70 

71 if image_rotary_emb is not None: 

72 query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) 

73 key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) 

74 

75 # Transpose for attention computation 

76 query = query.transpose(1, 2) 

77 key = key.transpose(1, 2) 

78 value = value.transpose(1, 2) 

79 

80 hidden_states = USP(query, key, value) 

81 

82 # Transpose back to original shape 

83 hidden_states = hidden_states.transpose(1, 2) 

84 

85 hidden_states = hidden_states.flatten(2, 3) 

86 hidden_states = hidden_states.to(query.dtype) 

87 

88 if encoder_hidden_states is not None: 

89 splits = hidden_states.split_with_sizes( 

90 [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 

91 ) 

92 encoder_hidden_states = splits[0] 

93 hidden_states = splits[1] 

94 encoder_hidden_states = attn.to_add_out(encoder_hidden_states) 

95 

96 hidden_states = attn.to_out[0](hidden_states) 

97 hidden_states = attn.to_out[1](hidden_states) 

98 

99 if encoder_hidden_states is not None: 

100 return hidden_states, encoder_hidden_states 

101 else: 

102 return hidden_states 

103 

104 

105@xFuserAttentionProcessorRegister.register(Flux2ParallelSelfAttnProcessor) 

106class xFuserFlux2ParallelSelfAttnProcessor(Flux2ParallelSelfAttnProcessor): 

107 

108 def __init__(self) -> None: 

109 super().__init__() 

110 

111 def __call__( 

112 self, 

113 attn: "Flux2ParallelSelfAttention", 

114 hidden_states: torch.Tensor, 

115 attention_mask: Optional[torch.Tensor] = None, 

116 image_rotary_emb: Optional[torch.Tensor] = None, 

117 ) -> torch.Tensor: 

118 # Parallel in (QKV + MLP in) projection 

119 hidden_states = attn.to_qkv_mlp_proj(hidden_states) 

120 qkv, mlp_hidden_states = torch.split( 

121 hidden_states, [3 * attn.inner_dim, attn.mlp_hidden_dim * attn.mlp_mult_factor], dim=-1 

122 ) 

123 

124 # Handle the attention logic 

125 query, key, value = qkv.chunk(3, dim=-1) 

126 

127 query = query.unflatten(-1, (attn.heads, -1)) 

128 key = key.unflatten(-1, (attn.heads, -1)) 

129 value = value.unflatten(-1, (attn.heads, -1)) 

130 

131 query = attn.norm_q(query) 

132 key = attn.norm_k(key) 

133 

134 if image_rotary_emb is not None: 

135 query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) # type: ignore[assignment] 

136 key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) # type: ignore[assignment] 

137 

138 # Transpose for attention computation 

139 query = query.transpose(1, 2) 

140 key = key.transpose(1, 2) 

141 value = value.transpose(1, 2) 

142 

143 hidden_states = USP(query, key, value) 

144 

145 # Transpose back to original shape 

146 hidden_states = hidden_states.transpose(1, 2) 

147 

148 hidden_states = hidden_states.flatten(2, 3) 

149 hidden_states = hidden_states.to(query.dtype) 

150 

151 # Handle the feedforward (FF) logic 

152 mlp_hidden_states = attn.mlp_act_fn(mlp_hidden_states) 

153 

154 # Concatenate and parallel output projection 

155 hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1) 

156 hidden_states = attn.to_out(hidden_states) 

157 

158 return hidden_states 

159 

160 

161@xFuserLayerWrappersRegister.register(Flux2ParallelSelfAttention) 

162class xFuserFlux2ParallelSelfAttention(xFuserAttentionBaseWrapper): 

163 

164 def __init__(self, attention: Flux2ParallelSelfAttention): 

165 super().__init__(attention=attention) 

166 self.processor = xFuserAttentionProcessorRegister.get_processor( 

167 attention.processor 

168 )() 

169 

170 def forward( 

171 self, 

172 hidden_states: torch.Tensor, 

173 attention_mask: Optional[torch.Tensor] = None, 

174 image_rotary_emb: Optional[torch.Tensor] = None, 

175 **kwargs: Any, 

176 ) -> torch.Tensor: 

177 

178 return super().forward( 

179 hidden_states, 

180 attention_mask, 

181 image_rotary_emb, 

182 **kwargs, 

183 ) 

184 

185 

186class xFuserFlux2Transformer2DWrapper(Flux2Transformer2DModel): 

187 

188 def __init__( 

189 self, 

190 patch_size: int = 1, 

191 in_channels: int = 128, 

192 out_channels: Optional[int] = None, 

193 num_layers: int = 8, 

194 num_single_layers: int = 48, 

195 attention_head_dim: int = 128, 

196 num_attention_heads: int = 48, 

197 joint_attention_dim: int = 15360, 

198 timestep_guidance_channels: int = 256, 

199 mlp_ratio: float = 3.0, 

200 axes_dims_rope: Tuple[int, ...] = (32, 32, 32, 32), 

201 rope_theta: int = 2000, 

202 eps: float = 1e-6, 

203 guidance_embeds: bool = True, 

204 ): 

205 super().__init__( 

206 patch_size=patch_size, 

207 in_channels=in_channels, 

208 out_channels=out_channels, 

209 num_layers=num_layers, 

210 num_single_layers=num_single_layers, 

211 attention_head_dim=attention_head_dim, 

212 num_attention_heads=num_attention_heads, 

213 joint_attention_dim=joint_attention_dim, 

214 timestep_guidance_channels=timestep_guidance_channels, 

215 mlp_ratio=mlp_ratio, 

216 axes_dims_rope=axes_dims_rope, 

217 rope_theta=rope_theta, 

218 eps=eps, 

219 guidance_embeds=guidance_embeds, 

220 ) 

221 

222 for block in self.transformer_blocks: 

223 block_any: Any = block 

224 block_any.attn.processor = xFuserFlux2AttnProcessor() 

225 for block in self.single_transformer_blocks: 

226 block_any = block 

227 block_any.attn.processor = xFuserFlux2ParallelSelfAttnProcessor() 

228 

229 def _pad_to_sp_divisible(self, tensor: torch.Tensor, padding_length: int, dim: int) -> torch.Tensor: 

230 padding = torch.zeros( 

231 *tensor.shape[:dim], padding_length, *tensor.shape[dim + 1:], dtype=tensor.dtype, device=tensor.device 

232 ) 

233 tensor = torch.cat([tensor, padding], dim=dim) 

234 return tensor 

235 

236 def forward( 

237 self, 

238 hidden_states: torch.Tensor, 

239 encoder_hidden_states: Optional[torch.Tensor] = None, 

240 *args: Any, 

241 timestep: Optional[torch.LongTensor] = None, 

242 img_ids: Optional[torch.Tensor] = None, 

243 txt_ids: Optional[torch.Tensor] = None, 

244 **kwargs: Any, 

245 ) -> Any: 

246 

247 sp_world_size = get_sequence_parallel_world_size() 

248 sequence_length = hidden_states.shape[1] 

249 padding_length = (sp_world_size - (sequence_length % sp_world_size)) % sp_world_size 

250 if padding_length > 0: 

251 hidden_states = self._pad_to_sp_divisible(hidden_states, padding_length, dim=1) 

252 assert img_ids is not None, "img_ids is required when padding" 

253 img_ids = self._pad_to_sp_divisible(img_ids, padding_length, dim=1) 

254 

255 if ( 

256 isinstance(timestep, torch.Tensor) 

257 and timestep.ndim != 0 

258 and timestep.shape[0] == hidden_states.shape[0] 

259 ): 

260 timestep = torch.chunk( 

261 timestep, get_classifier_free_guidance_world_size(), dim=0 

262 )[get_classifier_free_guidance_rank()] 

263 hidden_states = torch.chunk( 

264 hidden_states, get_classifier_free_guidance_world_size(), dim=0 

265 )[get_classifier_free_guidance_rank()] 

266 hidden_states = torch.chunk( 

267 hidden_states, get_sequence_parallel_world_size(), dim=-2 

268 )[get_sequence_parallel_rank()] 

269 assert encoder_hidden_states is not None, "encoder_hidden_states is required" 

270 assert img_ids is not None, "img_ids is required" 

271 assert txt_ids is not None, "txt_ids is required" 

272 encoder_hidden_states = torch.chunk( 

273 encoder_hidden_states, get_classifier_free_guidance_world_size(), dim=0 

274 )[get_classifier_free_guidance_rank()] 

275 encoder_hidden_states = torch.chunk( 

276 encoder_hidden_states, get_sequence_parallel_world_size(), dim=-2 

277 )[get_sequence_parallel_rank()] 

278 img_ids = torch.chunk(img_ids, get_sequence_parallel_world_size(), dim=-2)[get_sequence_parallel_rank()] 

279 txt_ids = torch.chunk(txt_ids, get_sequence_parallel_world_size(), dim=-2)[get_sequence_parallel_rank()] 

280 

281 output = super().forward( 

282 hidden_states, 

283 encoder_hidden_states, 

284 *args, 

285 timestep=timestep, 

286 img_ids=img_ids, 

287 txt_ids=txt_ids, 

288 **kwargs, 

289 ) 

290 

291 return_dict = not isinstance(output, tuple) 

292 sample = output[0] 

293 sample = get_sp_group().all_gather(sample, dim=-2) 

294 sample = get_cfg_group().all_gather(sample, dim=0) 

295 if padding_length > 0: 

296 sample = sample[:, :-padding_length, :] 

297 if return_dict: 

298 return output.__class__(sample, *output[1:]) 

299 return (sample, *output[1:])