Coverage for wrapper/wan/run_wan_benchmark.py: 57%

1import os 

2import time 

3import logging 

4import torch 

5 

6from PIL import Image 

7 

8from torch.profiler import profile 

9from torch.profiler import ProfilerActivity 

10 

11from wrapper_wan21 import Wan21VideoGeneration 

12 

13from wan.utils.utils import cache_video 

14 

15 

16def main() -> None: 

17 NUM_GPUS = 8 

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

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

20 GPU = torch.cuda.get_device_name(0) 

21 

22 TIMING_LOG_FILE = "timing.csv" 

23 

24 # Load the Wan model 

25 activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA, ProfilerActivity.XPU] 

26 with profile(activities=activities, profile_memory=True, record_shapes=True) as prof: 

27 t0 = time.time() 

28 

29 parallel_mode = "mixed" 

30 # Multiserver setup: ulysses in server and ring across 

31 ulysses_size = NUM_GPUS if WORLD_SIZE > NUM_GPUS else WORLD_SIZE 

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

33 

34 if parallel_mode == "ulysses": 

35 # Ulysses all GPUs in the cluster 

36 ulysses_size = WORLD_SIZE 

37 ring_size = 1 

38 elif parallel_mode == "ring": 

39 # Ring all GPUs in the cluster 

40 ulysses_size = 1 

41 ring_size = WORLD_SIZE 

42 

43 print(f"[{RANK:03d}] Parallel setup: {WORLD_SIZE} GPUs, {ulysses_size} Ulysses, {ring_size} Ring") 

44 video_gen = Wan21VideoGeneration( 

45 # ulysses_size=ulysses_size, 

46 # ring_size=ring_size, 

47 param_dtype=torch.bfloat16, 

48 # param_dtype=torch.float32, 

49 ) 

50 print(f"[{video_gen.rank}] Loaded model in {time.time() - t0:.3f} seconds {video_gen.load_timer}") 

51 

52 # Show profiling data 

53 if video_gen.rank == 0: 

54 print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10)) 

55 print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10)) 

56 

57 # Translated from "wan/configs/shared_config.py" 

58 input_neg_prompt = "Gorgeous colors, overexposed, static, blurred details, subtitles, style, artwork, " 

59 input_neg_prompt += "painting, picture, still, overall gray, worst quality, low quality, " 

60 input_neg_prompt += "JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, " 

61 input_neg_prompt += "poorly drawn faces, deformed, disfigured, deformed limbs, fused fingers, " 

62 input_neg_prompt += "still picture, cluttered background, three legs, many people in the background, " 

63 input_neg_prompt += "walking backwards" 

64 

65 # Query 1 

66 input_image = Image.open("generated_image.png").convert("RGB") # 640x480 

67 input_prompt = "The woman on the right speaks while the man on the left listens." 

68 input_neg_prompt = "No camera movement." 

69 

70 # Query 2 

71 # input_image = Image.open("generated_image_man.png").convert("RGB") 

72 # input_image = Image.open("generated_image_man_highres.png").convert("RGB") 

73 # input_prompt = "The man is explaining something to another person in the podcast studio." 

74 # input_neg_prompt = "No camera movement." 

75 

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

77 video = video_gen.generate( 

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

79 prompt="Warmup prompt", 

80 neg_prompt="Warmup prompt", 

81 num_frames=4 + 1, 

82 sampling_steps=1, 

83 ) 

84 

85 if video_gen.rank == 0: 

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

87 file_timing.write( 

88 "#run_id,hw,world_size,ulysses_size,ring_size,batch_size,num_frames,sampling_steps,colors," 

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

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

91 file_timing.write( 

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

93 

94 # Timing multiple configurations 

95 for run_id in range(1): 

96 # for num_frames in [80+1, 4+1, 8+1, 20+1, 40+1, 60+1, 80+1]: # 4n+1 

97 for num_frames in [80 + 1]: # 4n+1 

98 # for sampling_steps in [10, 1, 2, 4, 5, 6, 8, 10, 20, 50]: 

99 for sampling_steps in [10, 50]: 

100 t0 = time.time() 

101 video = video_gen.generate( 

102 img=input_image, 

103 prompt=input_prompt, 

104 neg_prompt=input_neg_prompt, 

105 num_frames=num_frames, 

106 sampling_steps=sampling_steps, 

107 ) 

108 total_time_seconds = time.time() - t0 

109 logging.info( 

110 f"Video generated in {total_time_seconds:.3f} seconds with {num_frames} frames " 

111 f"and {sampling_steps} steps.") 

112 

113 if video_gen.rank == 0: 

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

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

116 file_timing.write( 

117 f"{run_id},{GPU},{WORLD_SIZE},{ulysses_size},{ring_size},1,{num_frames}," 

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

119 

120 if video_gen.rank == 0: 

121 video_file_name = f"gen_video_world{WORLD_SIZE}_u{ulysses_size}_r{ring_size}_" + \ 

122 f"frames{num_frames}_steps{sampling_steps}_run{run_id}.mp4" 

123 cache_video( 

124 tensor=video[None], 

125 save_file=video_file_name, 

126 fps=video_gen.FPS, 

127 nrow=1, 

128 normalize=True, 

129 value_range=(-1, 1)) 

130 

131 del video_gen 

132 

133 

134if __name__ == "__main__": 

135 main()