Coverage for wrapper/flux/flux_xfuser.py: 89%

1# from https://github.com/xdit-project/xDiT/blob/main/examples/flux_usp_example.py 

2# from https://github.com/chengzeyi/ParaAttention/blob/main/examples/run_flux.py 

3 

4import torch 

5import functools 

6import types 

7 

8from typing import Optional 

9from typing import Any 

10 

11from diffusers import DiffusionPipeline 

12 

13from xfuser.core.distributed import get_runtime_state 

14from xfuser.core.distributed import get_classifier_free_guidance_world_size 

15from xfuser.core.distributed import get_classifier_free_guidance_rank 

16from xfuser.core.distributed import get_cfg_group 

17from xfuser.core.distributed import get_sequence_parallel_world_size 

18from xfuser.core.distributed import get_sequence_parallel_rank 

19from xfuser.core.distributed import get_sp_group 

20 

21from xfuser.model_executor.models.transformers.transformer_flux import xFuserFluxAttnProcessor 

22 

23 

24def parallelize_transformer(pipe: DiffusionPipeline) -> DiffusionPipeline: 

25 transformer = getattr(pipe, "transformer") 

26 assert transformer is not None, "pipe has no transformer attribute" 

27 original_forward = transformer.forward 

28 

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

30 def new_forward( 

31 self: Any, 

32 hidden_states: torch.Tensor, 

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

34 *args: Any, 

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

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

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

38 **kwargs: Any, 

39 ) -> Any: 

40 cfg_world = get_classifier_free_guidance_world_size() 

41 cfg_rank = get_classifier_free_guidance_rank() 

42 sp_world = get_sequence_parallel_world_size() 

43 sp_rank = get_sequence_parallel_rank() 

44 if hidden_states.shape[0] % cfg_world != 0: 

45 raise ValueError(f"Cannot split {hidden_states.shape[0]} ({hidden_states.shape}) into {cfg_world} parts") 

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

47 raise ValueError(f"Cannot split {hidden_states.shape[-2]} ({hidden_states.shape}) into {sp_world} parts") 

48 assert encoder_hidden_states is not None, "encoder_hidden_states must not be None" 

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

50 get_runtime_state().split_text_embed_in_sp = False 

51 else: 

52 get_runtime_state().split_text_embed_in_sp = True 

53 

54 if isinstance(timestep, torch.Tensor) and timestep.ndim != 0 and timestep.shape[0] == hidden_states.shape[0]: 

55 timestep = torch.chunk(timestep, cfg_world, dim=0)[cfg_rank] 

56 hidden_states = torch.chunk(hidden_states, cfg_world, dim=0)[cfg_rank] 

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

58 encoder_hidden_states = torch.chunk(encoder_hidden_states, cfg_world, dim=0)[cfg_rank] 

59 if get_runtime_state().split_text_embed_in_sp: 

60 encoder_hidden_states = torch.chunk(encoder_hidden_states, sp_world, dim=-2)[sp_rank] 

61 assert img_ids is not None, "img_ids must not be None" 

62 img_ids = torch.chunk(img_ids, sp_world, dim=-2)[sp_rank] 

63 if get_runtime_state().split_text_embed_in_sp: 

64 assert txt_ids is not None, "txt_ids must not be None when split_text_embed_in_sp is True" 

65 txt_ids = torch.chunk(txt_ids, sp_world, dim=-2)[sp_rank] 

66 

67 for block in transformer.transformer_blocks + transformer.single_transformer_blocks: 

68 block.attn.processor = xFuserFluxAttnProcessor() 

69 

70 output = original_forward( 

71 hidden_states, 

72 encoder_hidden_states, 

73 *args, 

74 timestep=timestep, 

75 img_ids=img_ids, 

76 txt_ids=txt_ids, 

77 **kwargs, 

78 ) 

79 

80 return_dict = not isinstance(output, tuple) 

81 sample = output[0] 

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

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

84 if return_dict: 

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

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

87 

88 bound_forward: Any = types.MethodType(new_forward, transformer) 

89 transformer.forward = bound_forward 

90 

91 return pipe