Coverage for wrapper/wan/run_wan_vae_decoder_benchmark.py: 0%
195 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2This script is to benchmark the VAE decoder and see how to optimize it.
3Trying to make the code parallel across 8 GPUs.
4"""
6from wan.utils.utils import cache_video
7from wan.modules.vae import WanVAE
9from typing import Any
10from typing import Iterator
11from typing import Optional
12from typing import Tuple
13from typing import Union
15import os
16import time
18import torch
19import torch.amp as amp
20import torch.distributed as dist
23ScalePair = Tuple[Union[float, torch.Tensor], Union[float, torch.Tensor]]
26# Original
27def decode_old(
28 vae_model: Any,
29 z: torch.Tensor,
30 scale: ScalePair
31) -> torch.Tensor:
32 vae_model.clear_cache()
33 # z: [b,c,t,h,w]
34 if isinstance(scale[0], torch.Tensor):
35 assert isinstance(scale[1], torch.Tensor)
36 z = z / scale[1].view(1, vae_model.z_dim, 1, 1, 1) + scale[0].view(
37 1, vae_model.z_dim, 1, 1, 1)
38 else:
39 z = z / scale[1] + scale[0]
40 iter_ = z.shape[2]
41 x = vae_model.conv2(z)
42 for i in range(iter_):
43 vae_model._conv_idx = [0]
44 if i == 0:
45 out = vae_model.decoder(
46 x[:, :, i:i + 1, :, :],
47 feat_cache=vae_model._feat_map,
48 feat_idx=vae_model._conv_idx)
49 else:
50 out_ = vae_model.decoder(
51 x[:, :, i:i + 1, :, :],
52 feat_cache=vae_model._feat_map,
53 feat_idx=vae_model._conv_idx)
54 out = torch.cat([out, out_], 2)
55 vae_model.clear_cache()
56 return out
59# New version of decode
60def decode(
61 vae_model: Any,
62 z: torch.Tensor,
63 scale: ScalePair,
64) -> torch.Tensor:
65 rank = dist.get_rank()
66 vae_model.clear_cache()
67 # z: [b,c,t,h,w]
68 if isinstance(scale[0], torch.Tensor):
69 assert isinstance(scale[1], torch.Tensor)
70 scale_0 = scale[0].view(1, vae_model.z_dim, 1, 1, 1)
71 scale_1 = scale[1].view(1, vae_model.z_dim, 1, 1, 1)
72 z = z / scale_1 + scale_0
73 else:
74 z = z / scale[1] + scale[0]
75 iter_ = z.shape[2]
76 x = vae_model.conv2(z)
77 out_list = []
78 for i in range(iter_):
79 vae_model._conv_idx = [0]
80 out_ = vae_model.decoder(
81 x[:, :, i:i + 1, :, :],
82 feat_cache=vae_model._feat_map,
83 feat_idx=vae_model._conv_idx)
84 if rank == 0:
85 print("DECODE", x.shape, out_.shape)
86 # 21 x torch.Size([1, 16, 1, 68, 90]) -> torch.Size([1, 3, 4, 544, 720])
87 out_list.append(out_)
88 out = torch.cat(out_list, dim=2)
89 vae_model.clear_cache()
90 return out
93# Distributed version
94# TODO need to remove the hard coding and clean it up
95def decode_parallel(
96 vae_model: Any,
97 z: torch.Tensor,
98 scale: ScalePair,
99) -> Optional[torch.Tensor]:
100 rank = dist.get_rank()
101 world_size = dist.get_world_size()
103 vae_model.clear_cache()
104 # z: [b,c,t,h,w]
105 if isinstance(scale[0], torch.Tensor):
106 assert isinstance(scale[1], torch.Tensor)
107 scale_0 = scale[0].view(1, vae_model.z_dim, 1, 1, 1)
108 scale_1 = scale[1].view(1, vae_model.z_dim, 1, 1, 1)
109 z = z / scale_1 + scale_0
110 else:
111 z = z / scale[1] + scale[0]
113 x = vae_model.conv2(z) # [b, c, t, h, w]
115 total_t = x.shape[2]
116 # slice_t = total_t // world_size
117 start = rank * total_t // world_size
118 start = max(0, start - 2) # We process the previous 2 frames for the convolution -> TODO take only the right output
119 end = (rank + 1) * total_t // world_size
120 local_x = x[:, :, start:end, :, :] # Local slice: [16, 21, 68, 90] -> [16, 2-3, 68, 90]
121 local_t = local_x.shape[2]
123 # Decode local slice
124 local_outs = []
125 for i in range(local_t):
126 vae_model._conv_idx = [0]
127 # TODO this cannot be split like this because it has dependencies
128 # [2] local_x:([1, 16, 2, 68, 90]) local_result:([1, 3, 5, 544, 720])
129 # [4] local_x:([1, 16, 3, 68, 90]) local_result:([1, 3, 9, 544, 720])
130 # conv -> middle -> upsamples -> head
131 out_ = vae_model.decoder(
132 local_x[:, :, i:i + 1, :, :],
133 feat_cache=vae_model._feat_map,
134 feat_idx=vae_model._conv_idx)
135 # print(f"[{rank}] DECODE each frame of", local_x.shape, "->", out_.shape)
136 # It should be [1, 16, 1(N), 68, 90] -> [1, 3, 4, 544, 720])
137 # but it does -> [1, 3, 1, 544, 720])
138 if rank == 0:
139 local_outs.append(out_)
140 elif i >= 2: # TODO this is a hack to avoid the first 2 frames
141 local_outs.append(out_)
142 local_result = torch.cat(local_outs, dim=2)
144 print(f"[{rank}] local_x:{local_x.shape} local_result: {local_result.shape}") # [b, c, slice_t, h, w]
145 # local_x:torch.Size([1, 16, 3, 68, 90]) local_result: torch.Size([1, 3, 3, 544, 720])
146 # local_x:torch.Size([1, 16, 2, 68, 90]) local_result: torch.Size([1, 3, 2, 544, 720])
148 # TODO The chunks have different sizes, so we need to gather them
150 device = None # TODO
151 # Extend tensor of size [1, 3, 5-12, 544, 720] to max size [1, 3, 12, 544, 720] with zeros
152 local_result_zeroes = torch.zeros(1, 3, 12, 544, 720).to(device) # [1, 3, 12, 544, 720]
153 local_result_zeroes[:, :, 0:local_result.shape[2], :, :] = local_result
154 local_result = local_result_zeroes
156 # Rank 0 gathers everything
157 # 1+2*10 -> 1+8*10
158 if rank == 0:
159 # gather_list = [torch.zeros_like(local_result) for _ in range(world_size)]
160 gather_list = [torch.zeros(1, 3, 12, 544, 720).to(device) for _ in range(world_size)] # TODO max size
161 else:
162 gather_list = None
163 dist.gather(local_result, gather_list=gather_list, dst=0) # This might be fast enough but needs warm up
165 final = None # TODO return something empty at least
166 if rank == 0:
167 assert gather_list is not None
168 # TODO avoid concatenating the 0s for padding
169 for gather_rank in range(len(gather_list)):
170 # TODO avoid the hardcoding
171 gather_size = 12
172 if gather_rank == 0:
173 gather_size = 5
174 elif gather_rank == 2 or gather_rank == 6:
175 gather_size = 8
176 gather_list[gather_rank] = gather_list[gather_rank][:, :, 0:gather_size, :, :] # TODO remove the 0s
178 final = torch.cat(gather_list, dim=2) # Concatenate on time axis
179 # [c, total_t, h, w] -> [3, 40!, 544, 720] Should be [3, 81, 544, 720]
180 print(f"[{rank}] final shape: {final.shape}")
182 vae_model.clear_cache()
183 return final
186def decode_stream(
187 vae_model: Any,
188 z: torch.Tensor,
189 scale: ScalePair,
190 start_frame: int = 0
191) -> Iterator[torch.Tensor]:
192 # TODO start_frame
193 vae_model.clear_cache()
194 # z: [b, c, #lat_frames, lat_h, lat_w]
195 if isinstance(scale[0], torch.Tensor):
196 assert isinstance(scale[1], torch.Tensor)
197 scale_0 = scale[0].view(1, vae_model.z_dim, 1, 1, 1)
198 scale_1 = scale[1].view(1, vae_model.z_dim, 1, 1, 1)
199 z = z / scale_1 + scale_0
200 else:
201 z = z / scale[1] + scale[0]
202 iter_ = z.shape[2]
203 x = vae_model.conv2(z)
204 for i in range(iter_):
205 vae_model._conv_idx = [0]
206 out = vae_model.decoder(
207 x[:, :, i:i + 1, :, :],
208 feat_cache=vae_model._feat_map,
209 feat_idx=vae_model._conv_idx)
210 num_frames = out.shape[2]
211 print("out shape", out.shape, num_frames)
212 for frame_ix in range(num_frames):
213 yield out[:, :, frame_ix, :, :].unsqueeze(2) # #frames x [b, c, h, w]
214 vae_model.clear_cache()
217def main() -> None:
218 # Setup distributed env
219 t0 = time.time()
220 rank = int(os.getenv("RANK", 0))
221 local_rank = int(os.getenv("LOCAL_RANK", 0))
222 world_size = int(os.getenv("WORLD_SIZE", 1))
223 print(f"[{rank}] Setting up distributed environment...")
225 device_id = local_rank
226 device = torch.device(f"cuda:{device_id}")
228 torch.cuda.set_device(local_rank)
230 dist.init_process_group(
231 backend="nccl",
232 init_method="env://",
233 rank=rank,
234 world_size=world_size,
235 )
236 print(f"[{rank}] Distributed environment setup in {time.time() - t0:.3f} seconds")
238 # Load VAE
239 t0 = time.time()
240 print(f"[{rank}] Loading WanVAE...")
241 ckpt_dir = "Wan2.1/Wan2.1-I2V-14B-480P"
242 vae = WanVAE(
243 vae_pth=os.path.join(ckpt_dir, 'Wan2.1_VAE.pth'),
244 device=device,
245 )
246 print(f"[{rank}] Loaded WanVAE in {time.time() - t0:.3f} seconds") # ~9 seconds
248 # 0 DEBUG: VAE decode z shape torch.Size([1, 16, 21, 68, 90])
249 x0 = torch.load("tensor_x0.pt", weights_only=True).to(device) # [16, 21, 68, 90]
250 print(f"[{rank}] x0 shape: {x0.shape}")
251 x0 = [x0]
253 # Warmup run (loading models fully, etc)
254 t0 = time.time()
255 print(f"[{rank}] Warmup decoding x0...")
256 videos = vae.decode(x0)
257 print(f"[{rank}] Decoded x0 in {time.time() - t0:.3f} seconds") # ~10.5 seconds
258 video = videos[0]
259 sum_val = video.sum().item()
260 print(f"[{rank}] video sum:", sum_val) # -33767824
261 print(f"[{rank}] videos shape: {video.shape}") # [3, 81, 544, 720]
263 # Sync just in case
264 dist.barrier()
266 # Actual run
267 t0 = time.time()
268 print(f"[{rank}] Decoding x0...")
269 videos = vae.decode(x0)
270 print(f"[{rank}] Decoded x0 in {time.time() - t0:.3f} seconds") # ~4.5 seconds
271 # Checking sizes
272 video = videos[0]
273 sum_val = video.sum().item()
274 print(f"[{rank}] video sum:", sum_val) # -33767824
275 print(f"[{rank}] videos shape: {video.shape}") # [3, 81, 544, 720]
277 cache_video(
278 tensor=video[None],
279 save_file="video_debug_original.mp4",
280 fps=16,
281 nrow=1,
282 )
284 # Inner code for decode distributed
285 for retry in range(1):
286 print(f"[{rank}] {retry} ==========================")
287 zs = x0
288 t0 = time.time()
289 with amp.autocast('cuda', dtype=vae.dtype):
290 '''
291 videos = [
292 # Replaced with the function
293 # vae.model.decode(u.unsqueeze(0), vae.scale).float().clamp_(-1, 1).squeeze(0)
294 decode(vae.model, u, vae.scale).float().clamp_(-1, 1).squeeze(0)
295 for u in zs
296 ]
297 '''
298 videos = []
299 for u in zs:
300 video = decode_parallel(vae.model, u, vae.scale)
301 # video = decode(vae.model, u, vae.scale)
302 if rank == 0:
303 assert video is not None
304 videos.append(video.float().clamp_(-1, 1).squeeze(0))
305 print(f"[{rank}] Parallelized inner decode in {time.time() - t0:.3f} seconds") # ~4.4 seconds
306 if rank == 0:
307 video = videos[0]
308 sum_val = video.sum().item()
309 print(f"[{rank}@{retry}] video sum:", sum_val) # -33351406.0 != -33767824
310 print(f"[{rank}@{retry}] videos shape: {video.shape}") # [3, 81, 544, 720]
311 # cache_video(
312 # tensor=video[None],
313 # save_file=f"video_debug_parallel_{retry}.mp4",
314 # fps=16,
315 # nrow=1)
317 # Inner code for decode yield
318 for retry in range(1):
319 print(f"[{rank}] {retry} ==========================")
320 zs = x0
321 t0 = time.time()
322 with amp.autocast('cuda', dtype=vae.dtype):
323 video_frames = []
324 frame_id = 0
325 for video_latent in zs:
326 # video_latent shape [16, 21, 68, 90]
327 for video_frame in vae.decode_stream(video_latent):
328 # video_frame shape [3, 544, 720]
329 if rank == 0:
330 video_frames.append(video_frame)
331 print(frame_id, "video_frame", video_frame.shape)
332 frame_id += 1
333 print(f"[{rank}] Inner decode with yield in {time.time() - t0:.3f} seconds") # ~4.4 seconds
334 if rank == 0:
335 video = torch.stack(video_frames) # 81 * [3, 544, 720] -> [81, 3, 544, 720]
336 video = video.permute(1, 0, 2, 3) # [81, 3, 544, 720] -> [3, 81, 544, 720]
337 sum_val = video.sum().item()
338 print(f"[{rank}@{retry}] video sum:", sum_val) # -33351406.0 != -33767824
339 print(f"[{rank}@{retry}] videos shape: {video.shape}") # [3, 81, 544, 720]
341 cache_video(
342 tensor=video[None],
343 save_file=f"video_debug_yield_{retry}.mp4",
344 fps=16,
345 nrow=1,
346 )
349if __name__ == "__main__":
350 main()