from lightx2v.attentions.distributed.partial_heads_attn.attn import partial_heads_attn def parallelize_hunyuan(hunyuan_model):
hunyuan_model.transformer_infer.parallel_attention = partial_heads_attn
