Coverage for wrapper/fantasytalking/fantasytalking_xfuser.py: 25%

1import logging 

2 

3from typing import Any 

4from typing import List 

5from typing import Union 

6from typing import Optional 

7from typing import Callable 

8 

9import torch 

10import torch.amp as amp 

11from torch import Tensor 

12from torch.nn import Module 

13 

14from xfuser.core.distributed import get_sequence_parallel_rank 

15from xfuser.core.distributed import get_sequence_parallel_world_size 

16from xfuser.core.distributed import get_sp_group 

17 

18from diffsynth.models.wan_video_dit import sinusoidal_embedding_1d 

19 

20 

21def usp_fantasytalking_forward( 

22 self: Module, 

23 x_list: List[Tensor], # [C, T, H, W] 

24 timestep: Tensor, # [B] 

25 context: List[Tensor], # [L, C] 

26 seq_len: Union[int, Tensor], # [B] or scalar 

27 clip_fea: Optional[Tensor] = None, 

28 y: Optional[List[Tensor]] = None, 

29 use_gradient_checkpointing: bool = False, 

30 audio_proj: Optional[Module] = None, 

31 audio_context_lens: Optional[Tensor] = None, # [B] length per example 

32 latents_num_frames: Optional[int] = None, 

33 audio_scale: float = 1.0, 

34 **kwargs: Any, 

35) -> Tensor: 

36 """ 

37 Copied from: 

38 https://github.com/Fantasy-AMAP/fantasy-talking/blob/main/diffsynth/models/wan_video_dit.py 

39 And adjusted based on: 

40 https://github.com/Wan-Video/Wan2.1/blob/main/wan/distributed/xdit_context_parallel.py 

41 x_list: A list of videos each with shape [C, T, H, W]. 

42 timestep: [B]. 

43 context: A list of text embeddings each with shape [L, C]. 

44 """ 

45 if self.model_type == "i2v": 

46 assert clip_fea is not None and y is not None 

47 # params 

48 device = x_list[0].device 

49 if self.freqs.device != device: 

50 self.freqs = self.freqs.to(device) 

51 

52 if y is not None: 

53 x_list = [torch.cat([u, v], dim=0) for u, v in zip(x_list, y)] 

54 

55 # embeddings 

56 x_embed_list = [ 

57 self.patch_embedding(u.unsqueeze(0)) # type: ignore[operator] 

58 for u in x_list 

59 ] 

60 grid_sizes = torch.stack( 

61 [ 

62 torch.tensor(u.shape[2:], dtype=torch.long) 

63 for u in x_embed_list 

64 ] 

65 ) # [B,2] 

66 x_t_list = [u.flatten(2).transpose(1, 2) for u in x_embed_list] # [[C, L, T],,] 

67 seq_lens = torch.tensor([u.size(1) for u in x_t_list], dtype=torch.long) 

68 assert seq_lens.max() <= seq_len 

69 x = torch.cat([ 

70 torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) 

71 for u in x_t_list 

72 ]) 

73 

74 # time embeddings 

75 with amp.autocast(dtype=torch.float32, device_type="cuda"): 

76 e = self.time_embedding( # type: ignore[operator] 

77 sinusoidal_embedding_1d(self.freq_dim, timestep).float() 

78 ) 

79 e0 = self.time_projection(e).unflatten(1, (6, self.dim)) # type: ignore[operator] 

80 assert e.dtype == torch.float32 and e0.dtype == torch.float32 

81 

82 # context 

83 context_lens = None 

84 context = self.text_embedding( # type: ignore[operator] 

85 torch.stack([ 

86 torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) # type: ignore[operator, arg-type] 

87 for u in context 

88 ]) 

89 ) 

90 

91 if clip_fea is not None: 

92 context_clip = self.img_emb(clip_fea) # type: ignore[operator] # bs x 257 x dim 

93 context = torch.concat([context_clip, context], dim=1) # type: ignore[assignment, list-item] 

94 

95 # Context Parallel 

96 sp_world = get_sequence_parallel_world_size() 

97 sp_rank = get_sequence_parallel_rank() 

98 if x.shape[1] % sp_world != 0: 

99 raise ValueError(f"Input sequence length {x.shape} is not divisible by sequence parallel {sp_world}") 

100 logging.debug(f"Input sequence length {x.shape} and sequence parallel {sp_world}.") 

101 x = torch.chunk(x, sp_world, dim=1)[sp_rank] 

102 logging.debug(f"Input sequence length after chunking {x.shape}.") 

103 

104 """ 

105 # https://github.com/Fantasy-AMAP/fantasy-talking/issues/52 

106 # audio chunking along the #frames dimension doesn't work -> black frames 

107 if audio_proj is not None: 

108 # chunking audio_proj based on sequence parallel rank 

109 logging.info(f"Audio projection shape {audio_proj.shape} and sequence parallel {sp_world}.") 

110 if audio_proj.shape[2] % sp_world != 0: 

111 # insert silence frames evenly to each chunk to make it divisible by sequence parallel 

112 # e.g. audio_proj.shape = [1, 4, 15, 2048] and sp_world = 4 -> [1, 4, 16, 2048] 

113 # calculate how many frames we need to pad 

114 num_frames = audio_proj.shape[2] 

115 pad_frames = math.ceil(num_frames / sp_world) * sp_world - num_frames 

116 

117 pad_shape = list(audio_proj.shape) 

118 pad_shape[2] = pad_frames 

119 silence = torch.zeros(pad_shape, dtype=audio_proj.dtype, device=audio_proj.device) 

120 # audio_proj = torch.cat([audio_proj, silence], dim=2) 

121 logging.info(f"Audio projection shape after padding {audio_proj.shape}.") 

122 audio_proj = torch.chunk(audio_proj, sp_world, dim=2)[sp_rank] 

123 else: 

124 audio_proj = torch.chunk(audio_proj, sp_world, dim=2)[sp_rank] 

125 logging.info(f"Audio projection shape after chunking {audio_proj.shape}.") 

126 """ 

127 

128 # arguments 

129 kwargs = dict( 

130 e=e0, 

131 seq_lens=seq_lens, 

132 grid_sizes=grid_sizes, 

133 freqs=self.freqs, 

134 context=context, 

135 context_lens=context_lens, 

136 audio_proj=audio_proj, 

137 audio_context_lens=audio_context_lens, 

138 latents_num_frames=latents_num_frames, 

139 audio_scale=audio_scale, 

140 ) 

141 

142 def create_custom_forward(module: Module) -> Callable[..., Tensor]: 

143 def custom_forward(*inputs: Tensor, **kwargs: Any) -> Tensor: 

144 return module(*inputs, **kwargs) 

145 return custom_forward 

146 

147 for block in self.blocks: # type: ignore[union-attr] 

148 if self.training and use_gradient_checkpointing: 

149 x = torch.utils.checkpoint.checkpoint( 

150 create_custom_forward(block), 

151 x, 

152 **kwargs, 

153 use_reentrant=False, 

154 ) 

155 else: 

156 x = block(x, **kwargs) 

157 

158 # head 

159 x = self.head(x, e) # type: ignore[operator] 

160 

161 # Context Parallel 

162 x = get_sp_group().all_gather(x, dim=1) 

163 

164 # unpatchify 

165 x = self.unpatchify(x, grid_sizes) # type: ignore[operator] 

166 x = torch.stack(x).float() 

167 return x