Coverage for wrapper/hunyuanframepack/hunyuanframepack_xfuser.py: 21%

1# mypy: ignore-errors 

2""" 

3xFuser FramePack for Hunyuan Video. 

4Extends the Hunyuan Video Framepack Pipeline to support xFuser's long context attention and sequence parallelism. 

5Sources: 

6* https://github.com/lllyasviel/FramePack 

7* https://github.com/xdit-project/xDiT 

8""" 

9import torch 

10import functools 

11import einops 

12 

13from typing import Tuple 

14from typing import Any 

15from typing import Optional 

16from typing import Dict 

17 

18from diffusers import HunyuanVideoFramepackPipeline 

19 

20from diffusers.models.attention import Attention 

21from diffusers.models.transformers.transformer_hunyuan_video import HunyuanVideoAttnProcessor2_0 

22from diffusers.models.transformers.transformer_2d import Transformer2DModelOutput 

23 

24from flash_attn import flash_attn_varlen_func 

25 

26from xfuser.envs import PACKAGES_CHECKER 

27from xfuser.core.cache_manager.cache_manager import get_cache_manager 

28from xfuser.core.distributed import get_sp_group 

29from xfuser.core.distributed import get_sequence_parallel_world_size 

30from xfuser.core.distributed import get_sequence_parallel_rank 

31from xfuser.core.distributed import get_classifier_free_guidance_world_size 

32from xfuser.core.distributed import get_classifier_free_guidance_rank 

33from xfuser.core.distributed import get_runtime_state 

34from xfuser.core.long_ctx_attention import xFuserLongContextAttention 

35from xfuser.model_executor.layers.attention_processor import xFuserAttentionProcessorRegister 

36 

37 

38# Class for the attention extending: 

39# https://github.com/xdit-project/xDiT/blob/main/xfuser/model_executor/layers/attention_processor.py 

40env_info = PACKAGES_CHECKER.get_packages_info() 

41HAS_LONG_CTX_ATTN = env_info["has_long_ctx_attn"] 

42HAS_FLASH_ATTN = env_info["has_flash_attn"] 

43 

44 

45@xFuserAttentionProcessorRegister.register(HunyuanVideoAttnProcessor2_0) 

46class xFuserFramepackSingleHunyuanVideoAttnProcessor2_0(HunyuanVideoAttnProcessor2_0): 

47 def __init__(self) -> None: 

48 super().__init__() 

49 

50 assert get_sequence_parallel_world_size() > 1 

51 assert HAS_LONG_CTX_ATTN is True 

52 assert HAS_FLASH_ATTN is True 

53 

54 self.hybrid_seq_parallel_attn = xFuserLongContextAttention(use_kv_cache=True) 

55 

56 # HunyuanAttnProcessorFlashAttnSingle 

57 def __call__( 

58 self, 

59 attn: HunyuanVideoAttnProcessor2_0, 

60 hidden_states: torch.Tensor, 

61 encoder_hidden_states: torch.Tensor, 

62 attention_mask: torch.Tensor, 

63 image_rotary_emb: torch.Tensor 

64 ) -> Tuple[torch.Tensor, torch.Tensor]: 

65 cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask 

66 hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) 

67 

68 query = attn.to_q(hidden_states) 

69 key = attn.to_k(hidden_states) 

70 value = attn.to_v(hidden_states) 

71 

72 query = query.unflatten(2, (attn.heads, -1)) 

73 key = key.unflatten(2, (attn.heads, -1)) 

74 value = value.unflatten(2, (attn.heads, -1)) 

75 

76 query = attn.norm_q(query) 

77 key = attn.norm_k(key) 

78 

79 txt_length = encoder_hidden_states.shape[1] 

80 query = torch.cat([apply_rotary_emb_transposed(query[:, :-txt_length], 

81 image_rotary_emb), query[:, -txt_length:]], dim=1) 

82 key = torch.cat([apply_rotary_emb_transposed(key[:, :-txt_length], 

83 image_rotary_emb), key[:, -txt_length:]], dim=1) 

84 hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, 

85 self.hybrid_seq_parallel_attn, attn) 

86 hidden_states = hidden_states.flatten(-2) 

87 hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:] 

88 return hidden_states, encoder_hidden_states 

89 

90 

91def attn_varlen_func( 

92 q: torch.Tensor, 

93 k: torch.Tensor, 

94 v: torch.Tensor, 

95 cu_seqlens_q: torch.Tensor, 

96 cu_seqlens_kv: torch.Tensor, 

97 max_seqlen_q: int, 

98 max_seqlen_kv: int, 

99 hybrid_seq_parallel_attn: Any = None, 

100 attn: Any = None 

101) -> torch.Tensor: 

102 if cu_seqlens_q is None and cu_seqlens_kv is None and max_seqlen_q is None and max_seqlen_kv is None: 

103 # Needed for XDiT 

104 x = hybrid_seq_parallel_attn(attn, q, k, v, dropout_p=0.0, causal=False, joint_strategy="none") 

105 # Original 

106 # x = flash_attn_func(q, k, v) 

107 return x 

108 

109 B, L, H, C = q.shape 

110 

111 q = q.flatten(0, 1) 

112 k = k.flatten(0, 1) 

113 v = v.flatten(0, 1) 

114 

115 x = flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) 

116 

117 x = x.unflatten(0, (B, L)) 

118 

119 return x 

120 

121 

122def apply_rotary_emb_transposed( 

123 x: torch.Tensor, 

124 freqs_cis: torch.Tensor 

125) -> torch.Tensor: 

126 # https://github.com/lllyasviel/FramePack/blob/c5d375661a2557383f0b8da9d11d14c23b0c4eaf/diffusers_helper/models/hunyuan_video_packed.py#L190 

127 cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) 

128 x_real, x_imag = x.unflatten(-1, (-1, 2)).unbind(-1) 

129 x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) 

130 out = x.float() * cos + x_rotated.float() * sin 

131 out = out.to(x) 

132 return out 

133 

134 

135@xFuserAttentionProcessorRegister.register(HunyuanVideoAttnProcessor2_0) 

136class xFuserFramepackDoubleHunyuanVideoAttnProcessor2_0(HunyuanVideoAttnProcessor2_0): 

137 def __init__(self) -> None: 

138 super().__init__() 

139 

140 assert get_sequence_parallel_world_size() > 1 

141 assert HAS_LONG_CTX_ATTN is True 

142 assert HAS_FLASH_ATTN is True 

143 

144 self.hybrid_seq_parallel_attn = xFuserLongContextAttention(use_kv_cache=True) 

145 

146 # HunyuanAttnProcessorFlashAttnDouble 

147 def __call__( 

148 self, 

149 attn: HunyuanVideoAttnProcessor2_0, 

150 hidden_states: torch.Tensor, 

151 encoder_hidden_states: torch.Tensor, 

152 attention_mask: torch.Tensor, 

153 image_rotary_emb: torch.Tensor 

154 ) -> Tuple[torch.Tensor, torch.Tensor]: 

155 cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = attention_mask 

156 query = attn.to_q(hidden_states) 

157 key = attn.to_k(hidden_states) 

158 value = attn.to_v(hidden_states) 

159 

160 query = query.unflatten(2, (attn.heads, -1)) 

161 key = key.unflatten(2, (attn.heads, -1)) 

162 value = value.unflatten(2, (attn.heads, -1)) 

163 

164 query = attn.norm_q(query) 

165 key = attn.norm_k(key) 

166 

167 query = apply_rotary_emb_transposed(query, image_rotary_emb) 

168 key = apply_rotary_emb_transposed(key, image_rotary_emb) 

169 

170 encoder_query = attn.add_q_proj(encoder_hidden_states) 

171 encoder_key = attn.add_k_proj(encoder_hidden_states) 

172 encoder_value = attn.add_v_proj(encoder_hidden_states) 

173 

174 encoder_query = encoder_query.unflatten(2, (attn.heads, -1)) 

175 encoder_key = encoder_key.unflatten(2, (attn.heads, -1)) 

176 encoder_value = encoder_value.unflatten(2, (attn.heads, -1)) 

177 

178 encoder_query = attn.norm_added_q(encoder_query) 

179 encoder_key = attn.norm_added_k(encoder_key) 

180 

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

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

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

184 hidden_states = attn_varlen_func(query, key, value, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, 

185 self.hybrid_seq_parallel_attn, attn) 

186 hidden_states = hidden_states.flatten(-2) 

187 txt_length = encoder_hidden_states.shape[1] 

188 hidden_states, encoder_hidden_states = hidden_states[:, :-txt_length], hidden_states[:, -txt_length:] 

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

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

191 encoder_hidden_states = attn.to_add_out(encoder_hidden_states) 

192 return hidden_states, encoder_hidden_states 

193 

194 

195def parallelize_transformer(pipe_hunyuan: HunyuanVideoFramepackPipeline) -> HunyuanVideoFramepackPipeline: 

196 transformer = pipe_hunyuan.transformer 

197 

198 """ 

199 Parallelize the transformer. 

200 """ 

201 @functools.wraps(transformer.__class__.forward) 

202 def new_forward( 

203 self: Any, 

204 hidden_states: torch.Tensor, # shape: [1, 16, 9, 80, 76] 

205 timestep: Any, 

206 encoder_hidden_states: torch.Tensor, # shape: [1, 512, 4096] 

207 encoder_attention_mask: torch.Tensor, # shape: [1, 512] 

208 pooled_projections: Any, 

209 guidance: Any, 

210 latent_indices: Any = None, 

211 clean_latents: Optional[torch.Tensor] = None, # shape: [1, 16, 2, 80, 76] 

212 clean_latent_indices: Optional[torch.Tensor] = None, # shape: [1, 2] 

213 clean_latents_2x: Optional[torch.Tensor] = None, # shape: [1, 16, 2, 80, 76] 

214 clean_latent_2x_indices: Optional[torch.Tensor] = None, # shape: [1, 2] 

215 clean_latents_4x: Optional[torch.Tensor] = None, # shape: [1, 16, 16, 80, 76] 

216 clean_latent_4x_indices: Optional[torch.Tensor] = None, # shape: [1, 16] 

217 image_embeddings: Optional[torch.Tensor] = None, 

218 attention_kwargs: Optional[Dict] = None, 

219 return_dict: bool = True 

220 ) -> Any: 

221 """ 

222 Started from https://github.com/lllyasviel/FramePack/blob/main/diffusers_helper/models/hunyuan_video_packed.py 

223 Aligned with https://github.com/xdit-project/xDiT/blob/main/examples/hunyuan_video_usp_example.py 

224 """ 

225 if attention_kwargs is None: 

226 attention_kwargs = {} 

227 

228 batch_size, num_channels, num_frames, height, width = hidden_states.shape 

229 

230 p, p_t = self.config['patch_size'], self.config['patch_size_t'] # 2, 1 

231 post_patch_num_frames = num_frames // p_t 

232 post_patch_height = height // p 

233 post_patch_width = width // p 

234 

235 original_context_length = post_patch_num_frames * post_patch_height * post_patch_width 

236 

237 # 1. FramePack + RoPE 

238 # image_rotary_emb ~= rope_freqs 

239 # hidden_states: [1, 16, 9, 80, 76] -> [1, 17500, 3072] = [1, 7*5*5*5*5*4, 3*2*512] 

240 hidden_states, rope_freqs = self.process_input_hidden_states( 

241 hidden_states, latent_indices, 

242 clean_latents, clean_latent_indices, 

243 clean_latents_2x, clean_latent_2x_indices, 

244 clean_latents_4x, clean_latent_4x_indices 

245 ) 

246 

247 # 2. Conditional embeddings 

248 temb = self.time_text_embed(timestep, guidance, pooled_projections) 

249 # [1, 512, 4096] -> [1, 512, 3072] 

250 encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) 

251 

252 extra_encoder_hidden_states = self.image_projection(image_embeddings) 

253 extra_attention_mask = torch.ones( 

254 (batch_size, extra_encoder_hidden_states.shape[1]), 

255 dtype=encoder_attention_mask.dtype, 

256 device=encoder_attention_mask.device 

257 ) 

258 encoder_hidden_states = torch.cat([extra_encoder_hidden_states, encoder_hidden_states], dim=1) 

259 encoder_attention_mask = torch.cat([extra_attention_mask, encoder_attention_mask], dim=1) 

260 

261 text_len = encoder_attention_mask.sum().item() 

262 # 1, 1241, 3072 -> 1, text_len(742), 3072 

263 encoder_hidden_states = encoder_hidden_states[:, :text_len] 

264 attention_mask = None, None, None, None 

265 

266 # Sequence-parallel chunking 

267 sp_world = get_sequence_parallel_world_size() 

268 sp_rank = get_sequence_parallel_rank() 

269 cfg_world = get_classifier_free_guidance_world_size() 

270 cfg_rank = get_classifier_free_guidance_rank() 

271 

272 # Chunk hidden_states and rope_freqs identically along the sequence dimension 

273 # [1, 17226, 3072] -> [1, 17226 / sp_world, 3072] 

274 # [1, 17500, 3072] -> [1, 17500 / sp_world, 3072] 

275 if hidden_states.shape[-2] % sp_world != 0: 

276 # This will cause a noisy output after chunking 

277 raise RuntimeError(f"hidden_states {hidden_states.shape} is not divisible by sp_world {sp_world}.") 

278 hidden_states = torch.chunk(hidden_states, sp_world, dim=-2)[sp_rank] 

279 

280 # [1, 17226, 256] -> [1, 17226 / sp_world, 256] 

281 if rope_freqs.shape[-2] % sp_world != 0: 

282 # This will cause a noisy output after chunking 

283 raise RuntimeError(f"rope_freqs {rope_freqs.shape} is not divisible by sp_world {sp_world}.") 

284 rope_freqs = torch.chunk(rope_freqs, sp_world, dim=-2)[sp_rank] 

285 

286 # Chunk encoder_hidden_states [1, 742, 3072] -> [1, 742 / sp_world, 3072] 

287 if encoder_hidden_states.shape[-2] % sp_world != 0: 

288 get_runtime_state().split_text_embed_in_sp = False 

289 else: 

290 get_runtime_state().split_text_embed_in_sp = True 

291 

292 encoder_hidden_states = torch.chunk( 

293 encoder_hidden_states, 

294 cfg_world, 

295 dim=0)[cfg_rank] 

296 

297 if get_runtime_state().split_text_embed_in_sp: 

298 encoder_hidden_states = torch.chunk( 

299 encoder_hidden_states, 

300 sp_world, 

301 dim=-2)[sp_rank] 

302 

303 # 3. Transformer blocks 

304 # https://github.com/lllyasviel/FramePack/blob/main/diffusers_helper/models/hunyuan_video_packed.py#HunyuanVideoTransformerBlock 

305 for block in self.transformer_blocks + self.single_transformer_blocks: 

306 # HunyuanVideoSingleTransformerBlock 

307 hidden_states, encoder_hidden_states = block( 

308 hidden_states, 

309 encoder_hidden_states, 

310 temb, 

311 attention_mask, 

312 rope_freqs 

313 ) 

314 

315 # Output projection 

316 # This is torch.Size([1, 8613, 3072]) 

317 hidden_states = self.norm_out(hidden_states, temb) 

318 # Ensure dtype matches proj_out weights 

319 hidden_states = hidden_states.to(self.proj_out.weight.dtype) 

320 hidden_states = self.proj_out(hidden_states) 

321 # shape: ([1, 8613, 64]) 

322 # Gather and reshape like xDiT 

323 hidden_states = get_sp_group().all_gather(hidden_states, dim=-2) 

324 hidden_states = hidden_states[:, -original_context_length:, :] 

325 hidden_states = einops.rearrange(hidden_states, 'b (t h w) (c pt ph pw) -> b c (t pt) (h ph) (w pw)', 

326 t=post_patch_num_frames, h=post_patch_height, w=post_patch_width, 

327 pt=p_t, ph=p, pw=p) 

328 # Rearranged hidden_states [1, 16, 9, 80, 76] 

329 

330 if return_dict: 

331 return Transformer2DModelOutput(sample=hidden_states) 

332 

333 return hidden_states, 

334 

335 new_forward = new_forward.__get__(transformer) # type: ignore[attr-defined] 

336 transformer.forward = new_forward 

337 

338 # Apply the xFuser attention processor to all transformer blocks 

339 for block in transformer.transformer_blocks: 

340 block.attn.processor = xFuserFramepackDoubleHunyuanVideoAttnProcessor2_0() 

341 for block in transformer.single_transformer_blocks: 

342 block.attn.processor = xFuserFramepackSingleHunyuanVideoAttnProcessor2_0() 

343 

344 # Register the attention layers to the cache manager 

345 # This is needed for sequence parallelism to work correctly and allow each chunk to attend to the entire sequence 

346 for block in transformer.transformer_blocks: 

347 for submodule in block.modules(): 

348 if isinstance(submodule, Attention): 

349 get_cache_manager().register_cache_entry(submodule, "attn", "sequence_parallel_attn_cache") 

350 for block in transformer.single_transformer_blocks: 

351 for submodule in block.modules(): 

352 if isinstance(submodule, Attention): 

353 get_cache_manager().register_cache_entry(submodule, "attn", "sequence_parallel_attn_cache") 

354 

355 return pipe_hunyuan