Coverage for wrapper/wan/run_wan_benchmark_batching.py: 35%

1import os 

2import torch 

3 

4from typing import List 

5 

6from PIL import Image 

7from PIL.Image import Resampling 

8 

9from wrapper_wan21 import Wan21VideoGeneration 

10 

11from wan.utils.utils import cache_video 

12 

13 

14def main() -> None: 

15 NUM_GPUS = 8 

16 WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1)) 

17 # RANK = int(os.environ.get("RANK", 1)) 

18 GPU = torch.cuda.get_device_name(0) 

19 

20 DEBUG = False 

21 

22 ulysses_size = NUM_GPUS if WORLD_SIZE > NUM_GPUS else WORLD_SIZE 

23 ring_size = WORLD_SIZE // NUM_GPUS if WORLD_SIZE > NUM_GPUS else 1 

24 

25 video_gen = Wan21VideoGeneration( 

26 # ulysses_size=ulysses_size, 

27 # ring_size=ring_size, 

28 ) 

29 

30 # Size of the frame taken by Wan 2.1 480p 

31 height = 544 

32 width = 720 

33 

34 input_images: List[Image.Image] = [ 

35 Image.open("generated_image_20250415T165936.png"), 

36 Image.open("generated_image_20250416T162934.png"), 

37 Image.open("generated_image_flux_20250416T190517.png"), 

38 Image.open("generated_image_hidream_20250416T171929.png"), 

39 Image.open("generated_image_hidream_20250416T172445.png"), 

40 Image.open("generated_image_hidream_20250416T190419.png"), 

41 Image.open("generated_image_hidream_20250416T174426.png"), 

42 Image.open("person1_generated_image_hidream_20250416T171929.png"), 

43 Image.open("person0_generated_image_hidream_20250416T171929.png"), 

44 ] 

45 input_images = [ 

46 input_image.convert("RGB").resize((width, height), Resampling.LANCZOS) 

47 for input_image in input_images 

48 ] 

49 

50 input_prompts = [ 

51 "The person is speaking.", 

52 "The person is covering their face.", 

53 "The person is raising their arms.", 

54 "The person is yawning.", 

55 "The person is standing up.", 

56 "The person is pointing.", 

57 "The person is sitting down.", 

58 "The person is waving.", 

59 "The person is clapping.", 

60 "The person is turning their head.", 

61 ] 

62 

63 # Warm-up run (loading models fully, etc) 

64 video = video_gen.generate( 

65 img=Image.open("warmup_image.png").convert("RGB"), 

66 prompt="Warmup prompt", 

67 neg_prompt="Warmup prompt", 

68 num_frames=1 + 4, 

69 sampling_steps=1, 

70 ) 

71 

72 if video_gen.rank == 0: 

73 TIMING_LOG_FILE = "timing.csv" 

74 with open(TIMING_LOG_FILE, "a") as file_timing: 

75 file_timing.write( 

76 "#run_id,hw,world_size,ulysses_size,ring_size,batch_size,num_frames,sampling_steps," 

77 "colors,frames,height,width,txt_enc,img_enc,vae_enc,sched_setup,dit,sched,vae_dec,total\n") 

78 video_size_str = ",".join(map(str, video.size())) 

79 file_timing.write( 

80 f"w,{GPU},{WORLD_SIZE},{ulysses_size},{ring_size},1,5,1,{video_size_str},{video_gen.gen_timer[-1]}\n") 

81 

82 # Run without batching to set the baseline 

83 video = video_gen.generate( 

84 img=input_images[0], 

85 prompt=input_prompts[0], 

86 neg_prompt="", 

87 num_frames=81, 

88 sampling_steps=10, 

89 ) 

90 if video_gen.rank == 0: 

91 with open(TIMING_LOG_FILE, "a") as file_timing: 

92 video_size_str = ",".join(map(str, video.size())) 

93 file_timing.write( 

94 f"b,{GPU},{WORLD_SIZE},{ulysses_size},{ring_size},1,81,10,{video_size_str},{video_gen.gen_timer[-1]}\n") 

95 

96 if DEBUG: 

97 video_file_name = f"gen_video_world{WORLD_SIZE}_u{ulysses_size}_r{ring_size}_frames81_steps10_run0.mp4" 

98 cache_video( 

99 tensor=video[None], 

100 save_file=video_file_name, 

101 fps=video_gen.FPS, 

102 nrow=1, 

103 normalize=True, 

104 value_range=(-1, 1) 

105 ) 

106 

107 # Running batches of different sizes 

108 for run_id in range(1): 

109 num_frames = 1 + 80 

110 sampling_steps = 10 

111 for batch_size in range(1, len(input_images) + 1): 

112 batch_img = input_images[0:batch_size] 

113 batch_prompt = input_prompts[0:batch_size] 

114 batch_neg_prompt = [""] * batch_size 

115 batch_num_frames = [num_frames] * batch_size 

116 batch_start_frames = [1] * batch_size 

117 

118 videos = video_gen.generate_batch( 

119 batch_img=batch_img, 

120 batch_prompt=batch_prompt, 

121 batch_neg_prompt=batch_neg_prompt, 

122 batch_num_frames=batch_num_frames, 

123 batch_start_frames=batch_start_frames, 

124 sampling_steps=sampling_steps, 

125 ) 

126 if video_gen.rank == 0: 

127 with open(TIMING_LOG_FILE, "a") as file_timing: 

128 video_size_str = ",".join(map(str, videos[0].size())) 

129 file_timing.write( 

130 f"{run_id},{GPU},{WORLD_SIZE},{ulysses_size},{ring_size},{batch_size},{num_frames}," 

131 f"{sampling_steps},{video_size_str},{video_gen.gen_timer[-1]}\n") 

132 

133 # Save videos for debugging 

134 if DEBUG: 

135 for batch_id, video in enumerate(videos): 

136 video_file_name = f"gen_video_world{WORLD_SIZE}_u{ulysses_size}_" + \ 

137 f"r{ring_size}_batch{batch_id}_frames{num_frames}_steps{sampling_steps}_run{run_id}.mp4" 

138 cache_video( 

139 tensor=video[None], 

140 save_file=video_file_name, 

141 fps=video_gen.FPS, 

142 nrow=1, 

143 normalize=True, 

144 value_range=(-1, 1) 

145 ) 

146 

147 

148if __name__ == "__main__": 

149 main()