Coverage for wrapper/wan/vae.py: 17%

1# mypy: ignore-errors

2# Copy from: https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/vae.py

3# Added some changes for streaming encode/decode and 1-frame encode/decode.

6import logging

8from typing import List

9from typing import Optional

10from typing import Tuple

11from typing import Generator

12from typing import Any

14import torch

15import torch.amp as amp

16import torch.nn as nn

17import torch.nn.functional as F

18from einops import rearrange

20__all__ = [

21 'WanVAE',

22]

24CACHE_T = 2

27class CausalConv3d(nn.Conv3d):

28 """

29 Causal 3d convolusion.

30 """

32 def __init__(

33 self,

34 *args: Any,

35 **kwargs: Any

36 ) -> None:

37 super().__init__(*args, **kwargs)

38 self._padding: Tuple[int, int, int, int, int] = (

39 self.padding[2],

40 self.padding[2],

41 self.padding[1],

42 self.padding[1],

43 2 * self.padding[0], 0

44 )

45 self.padding = (0, 0, 0)

47 def forward(

48 self,

49 x: torch.Tensor,

50 cache_x: Optional[torch.Tensor] = None,

51 ) -> torch.Tensor:

52 padding = list(self._padding)

53 if cache_x is not None and self._padding[4] > 0:

54 cache_x = cache_x.to(x.device)

55 x = torch.cat([cache_x, x], dim=2)

56 padding[4] -= cache_x.shape[2]

57 x = F.pad(x, padding)

59 return super().forward(x)

62class RMS_norm(nn.Module):

64 def __init__(

65 self,

66 dim: int,

67 channel_first: bool = True,

68 images: bool = True,

69 bias: bool = False

70 ) -> None:

71 super().__init__()

72 broadcastable_dims = (1, 1, 1) if not images else (1, 1)

73 shape = (dim, *broadcastable_dims) if channel_first else (dim,)

75 self.channel_first = channel_first

76 self.scale = dim**0.5

77 self.gamma = nn.Parameter(torch.ones(shape))

78 self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.

80 def forward(

81 self,

82 x: torch.Tensor,

83 ) -> torch.Tensor:

84 return F.normalize(

85 x,

86 dim=(1 if self.channel_first else -1)

87 ) * self.scale * self.gamma + self.bias

90class Upsample(nn.Upsample):

92 def forward(

93 self,

94 x: torch.Tensor,

95 ) -> torch.Tensor:

96 """

97 Fix bfloat16 support for nearest neighbor interpolation.

98 """

99 return super().forward(x.float()).type_as(x)

100

101

102class Resample(nn.Module):

103

104 def __init__(

105 self,

106 dim: int,

107 mode: str

108 ) -> None:

109 assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',

110 'downsample3d')

111 super().__init__()

112 self.dim = dim

113 self.mode = mode

114

115 # layers

116 if mode == 'upsample2d':

117 self.resample = nn.Sequential(

118 Upsample(scale_factor=(2., 2.), mode='nearest-exact'),

119 nn.Conv2d(dim, dim // 2, 3, padding=1))

120 elif mode == 'upsample3d':

121 self.resample = nn.Sequential(

122 Upsample(scale_factor=(2., 2.), mode='nearest-exact'),

123 nn.Conv2d(dim, dim // 2, 3, padding=1))

124 self.time_conv = CausalConv3d(

125 dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))

126

127 elif mode == 'downsample2d':

128 self.resample = nn.Sequential(

129 nn.ZeroPad2d((0, 1, 0, 1)),

130 nn.Conv2d(dim, dim, 3, stride=(2, 2)))

131 elif mode == 'downsample3d':

132 self.resample = nn.Sequential(

133 nn.ZeroPad2d((0, 1, 0, 1)),

134 nn.Conv2d(dim, dim, 3, stride=(2, 2)))

135 self.time_conv = CausalConv3d(

136 dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))

137 else:

138 self.resample = nn.Identity()

139

140 def forward(

141 self,

142 x: torch.Tensor,

143 feat_cache: Optional[List[Any]] = None,

144 feat_idx: List[int] = [0]

145 ) -> torch.Tensor:

146 b, c, t, h, w = x.size()

147 if self.mode == 'upsample3d':

148 if feat_cache is not None:

149 idx = feat_idx[0]

150 if feat_cache[idx] is None:

151 feat_cache[idx] = 'Rep'

152 feat_idx[0] += 1

153 else:

154

155 cache_x = x[:, :, -CACHE_T:, :, :].clone()

156 if cache_x.shape[2] < 2 and feat_cache[

157 idx] is not None and feat_cache[idx] != 'Rep':

158 # cache last frame of last two chunk

159 cache_x = torch.cat([

160 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),

161 cache_x

162 ],

163 dim=2)

164 if cache_x.shape[2] < 2 and feat_cache[

165 idx] is not None and feat_cache[idx] == 'Rep':

166 cache_x = torch.cat([

167 torch.zeros_like(cache_x).to(cache_x.device),

168 cache_x

169 ],

170 dim=2)

171 if feat_cache[idx] == 'Rep':

172 x = self.time_conv(x)

173 else:

174 x = self.time_conv(x, feat_cache[idx])

175 feat_cache[idx] = cache_x

176 feat_idx[0] += 1

177

178 x = x.reshape(b, 2, c, t, h, w)

179 x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),

180 3)

181 x = x.reshape(b, c, t * 2, h, w)

182 t = x.shape[2]

183 x = rearrange(x, 'b c t h w -> (b t) c h w')

184 x = self.resample(x)

185 x = rearrange(x, '(b t) c h w -> b c t h w', t=t)

186

187 if self.mode == 'downsample3d':

188 if feat_cache is not None:

189 idx = feat_idx[0]

190 if feat_cache[idx] is None:

191 feat_cache[idx] = x.clone()

192 feat_idx[0] += 1

193 else:

194

195 cache_x = x[:, :, -1:, :, :].clone()

196 # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':

197 # # cache last frame of last two chunk

198 # cache_x = torch.cat(

199 # [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),cache_x], dim=2)

200

201 x = self.time_conv(

202 torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))

203 feat_cache[idx] = cache_x

204 feat_idx[0] += 1

205 return x

206

207 def init_weight(

208 self,

209 conv: nn.Conv3d

210 ) -> None:

211 conv_weight = conv.weight

212 nn.init.zeros_(conv_weight)

213 c1, c2, t, h, w = conv_weight.size()

214 one_matrix = torch.eye(c1, c2)

215 init_matrix = one_matrix

216 nn.init.zeros_(conv_weight)

217 # conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5

218 conv_weight.data[:, :, 1, 0, 0] = init_matrix # * 0.5

219 conv.weight.data.copy_(conv_weight)

220 nn.init.zeros_(conv.bias.data)

221

222 def init_weight2(

223 self,

224 conv: nn.Conv3d

225 ) -> None:

226 conv_weight = conv.weight.data

227 nn.init.zeros_(conv_weight)

228 c1, c2, t, h, w = conv_weight.size()

229 init_matrix = torch.eye(c1 // 2, c2)

230 # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)

231 conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix

232 conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix

233 conv.weight.data.copy_(conv_weight)

234 nn.init.zeros_(conv.bias.data)

235

236

237class ResidualBlock(nn.Module):

238

239 def __init__(

240 self,

241 in_dim: int,

242 out_dim: int,

243 dropout: float = 0.0

244 ) -> None:

245 super().__init__()

246 self.in_dim = in_dim

247 self.out_dim = out_dim

248

249 # layers

250 self.residual = nn.Sequential(

251 RMS_norm(in_dim, images=False), nn.SiLU(),

252 CausalConv3d(in_dim, out_dim, 3, padding=1),

253 RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),

254 CausalConv3d(out_dim, out_dim, 3, padding=1))

255 self.shortcut = CausalConv3d(in_dim, out_dim, 1) \

256 if in_dim != out_dim else nn.Identity()

257

258 def forward(

259 self,

260 x: torch.Tensor,

261 feat_cache: Optional[List[Any]] = None,

262 feat_idx: List[int] = [0]

263 ) -> torch.Tensor:

264 h = self.shortcut(x)

265 for layer in self.residual:

266 if isinstance(layer, CausalConv3d) and feat_cache is not None:

267 idx = feat_idx[0]

268 cache_x = x[:, :, -CACHE_T:, :, :].clone()

269 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:

270 # cache last frame of last two chunk

271 cache_x = torch.cat([

272 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(

273 cache_x.device), cache_x

274 ],

275 dim=2)

276 x = layer(x, feat_cache[idx])

277 feat_cache[idx] = cache_x

278 feat_idx[0] += 1

279 else:

280 x = layer(x)

281 return x + h

282

283

284class AttentionBlock(nn.Module):

285 """

286 Causal self-attention with a single head.

287 """

288

289 def __init__(

290 self,

291 dim: int,

292 ) -> None:

293 super().__init__()

294 self.dim = dim

295

296 # layers

297 self.norm = RMS_norm(dim)

298 self.to_qkv = nn.Conv2d(dim, dim * 3, 1)

299 self.proj = nn.Conv2d(dim, dim, 1)

300

301 # zero out the last layer params

302 nn.init.zeros_(self.proj.weight)

303

304 def forward(

305 self,

306 x: torch.Tensor

307 ) -> torch.Tensor:

308 identity = x

309 b, c, t, h, w = x.size()

310 x = rearrange(x, 'b c t h w -> (b t) c h w')

311 x = self.norm(x)

312 # compute query, key, value

313 q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3,

314 -1).permute(0, 1, 3,

315 2).contiguous().chunk(

316 3, dim=-1)

317

318 # apply attention

319 x = F.scaled_dot_product_attention(

320 q,

321 k,

322 v,

323 )

324 x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)

325

326 # output

327 x = self.proj(x)

328 x = rearrange(x, '(b t) c h w-> b c t h w', t=t)

329 return x + identity

330

331

332class Encoder3d(nn.Module):

333

334 def __init__(

335 self,

336 dim: int = 128,

337 z_dim: int = 4,

338 dim_mult: List[int] = [1, 2, 4, 4],

339 num_res_blocks: int = 2,

340 attn_scales: List[float] = [],

341 temperal_downsample: List[bool] = [True, True, False],

342 dropout: float = 0.0

343 ) -> None:

344 super().__init__()

345 self.dim = dim

346 self.z_dim = z_dim

347 self.dim_mult = dim_mult

348 self.num_res_blocks = num_res_blocks

349 self.attn_scales = attn_scales

350 self.temperal_downsample = temperal_downsample

351

352 # dimensions

353 dims = [dim * u for u in [1] + dim_mult]

354 scale = 1.0

355

356 # init block

357 self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)

358

359 # downsample blocks

360 downsamples = []

361 for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):

362 # residual (+attention) blocks

363 for _ in range(num_res_blocks):

364 downsamples.append(ResidualBlock(in_dim, out_dim, dropout))

365 if scale in attn_scales:

366 downsamples.append(AttentionBlock(out_dim))

367 in_dim = out_dim

368

369 # downsample block

370 if i != len(dim_mult) - 1:

371 mode = 'downsample3d' if temperal_downsample[

372 i] else 'downsample2d'

373 downsamples.append(Resample(out_dim, mode=mode))

374 scale /= 2.0

375 self.downsamples = nn.Sequential(*downsamples)

376

377 # middle blocks

378 self.middle = nn.Sequential(

379 ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim),

380 ResidualBlock(out_dim, out_dim, dropout))

381

382 # output blocks

383 self.head = nn.Sequential(

384 RMS_norm(out_dim, images=False), nn.SiLU(),

385 CausalConv3d(out_dim, z_dim, 3, padding=1))

386

387 def forward(

388 self,

389 x: torch.Tensor,

390 feat_cache: Optional[List[torch.Tensor]] = None,

391 feat_idx: List[int] = [0]

392 ) -> torch.Tensor:

393 if feat_cache is not None:

394 idx = feat_idx[0]

395 cache_x = x[:, :, -CACHE_T:, :, :].clone()

396 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:

397 # cache last frame of last two chunk

398 cache_x = torch.cat([

399 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(

400 cache_x.device), cache_x

401 ],

402 dim=2)

403 x = self.conv1(x, feat_cache[idx])

404 feat_cache[idx] = cache_x

405 feat_idx[0] += 1

406 else:

407 x = self.conv1(x)

408

409 # downsamples

410 for layer in self.downsamples:

411 if feat_cache is not None:

412 x = layer(x, feat_cache, feat_idx)

413 else:

414 x = layer(x)

415

416 # middle

417 for layer in self.middle:

418 if isinstance(layer, ResidualBlock) and feat_cache is not None:

419 x = layer(x, feat_cache, feat_idx)

420 else:

421 x = layer(x)

422

423 # head

424 for layer in self.head:

425 if isinstance(layer, CausalConv3d) and feat_cache is not None:

426 idx = feat_idx[0]

427 cache_x = x[:, :, -CACHE_T:, :, :].clone()

428 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:

429 # cache last frame of last two chunk

430 cache_x = torch.cat([

431 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(

432 cache_x.device), cache_x

433 ],

434 dim=2)

435 x = layer(x, feat_cache[idx])

436 feat_cache[idx] = cache_x

437 feat_idx[0] += 1

438 else:

439 x = layer(x)

440 return x

441

442

443class Decoder3d(nn.Module):

444

445 def __init__(

446 self,

447 dim: int = 128,

448 z_dim: int = 4,

449 dim_mult: List[int] = [1, 2, 4, 4],

450 num_res_blocks: int = 2,

451 attn_scales: List[float] = [],

452 temperal_upsample: List[bool] = [False, True, True],

453 dropout: float = 0.0

454 ) -> None:

455 super().__init__()

456 self.dim = dim

457 self.z_dim = z_dim

458 self.dim_mult = dim_mult

459 self.num_res_blocks = num_res_blocks

460 self.attn_scales = attn_scales

461 self.temperal_upsample = temperal_upsample

462

463 # dimensions

464 dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]

465 scale = 1.0 / 2**(len(dim_mult) - 2)

466

467 # init block

468 self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)

469

470 # middle blocks

471 self.middle = nn.Sequential(

472 ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),

473 ResidualBlock(dims[0], dims[0], dropout))

474

475 # upsample blocks

476 upsamples = []

477 for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):

478 # residual (+attention) blocks

479 if i == 1 or i == 2 or i == 3:

480 in_dim = in_dim // 2

481 for _ in range(num_res_blocks + 1):

482 upsamples.append(ResidualBlock(in_dim, out_dim, dropout))

483 if scale in attn_scales:

484 upsamples.append(AttentionBlock(out_dim))

485 in_dim = out_dim

486

487 # upsample block

488 if i != len(dim_mult) - 1:

489 mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'

490 upsamples.append(Resample(out_dim, mode=mode))

491 scale *= 2.0

492 self.upsamples = nn.Sequential(*upsamples)

493

494 # output blocks

495 self.head = nn.Sequential(

496 RMS_norm(out_dim, images=False), nn.SiLU(),

497 CausalConv3d(out_dim, 3, 3, padding=1))

498

499 def forward(

500 self,

501 x: torch.Tensor,

502 feat_cache: Optional[List[torch.Tensor]] = None,

503 feat_idx: List[int] = [0]

504 ) -> torch.Tensor:

505 # conv1

506 if feat_cache is not None:

507 idx = feat_idx[0]

508 cache_x = x[:, :, -CACHE_T:, :, :].clone()

509 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:

510 # cache last frame of last two chunk

511 cache_x = torch.cat([

512 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(

513 cache_x.device), cache_x

514 ],

515 dim=2)

516 x = self.conv1(x, feat_cache[idx])

517 feat_cache[idx] = cache_x

518 feat_idx[0] += 1

519 else:

520 x = self.conv1(x)

521

522 # middle

523 for layer in self.middle:

524 if isinstance(layer, ResidualBlock) and feat_cache is not None:

525 x = layer(x, feat_cache, feat_idx)

526 else:

527 x = layer(x)

528

529 # upsamples

530 for layer in self.upsamples:

531 if feat_cache is not None:

532 x = layer(x, feat_cache, feat_idx)

533 else:

534 x = layer(x)

535

536 # head

537 for layer in self.head:

538 if isinstance(layer, CausalConv3d) and feat_cache is not None:

539 idx = feat_idx[0]

540 cache_x = x[:, :, -CACHE_T:, :, :].clone()

541 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:

542 # cache last frame of last two chunk

543 cache_x = torch.cat([

544 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(

545 cache_x.device), cache_x

546 ],

547 dim=2)

548 x = layer(x, feat_cache[idx])

549 feat_cache[idx] = cache_x

550 feat_idx[0] += 1

551 else:

552 x = layer(x)

553 return x

554

555

556def count_conv3d(model: nn.Module) -> int:

557 count = 0

558 for m in model.modules():

559 if isinstance(m, CausalConv3d):

560 count += 1

561 return count

562

563

564class WanVAE_(nn.Module):

565

566 def __init__(

567 self,

568 dim: int = 128,

569 z_dim: int = 4,

570 dim_mult: List[int] = [1, 2, 4, 4],

571 num_res_blocks: int = 2,

572 attn_scales: List[float] = [],

573 temperal_downsample: List[bool] = [True, True, False],

574 dropout: float = 0.0

575 ) -> None:

576 super().__init__()

577 self.dim = dim

578 self.z_dim = z_dim

579 self.dim_mult = dim_mult

580 self.num_res_blocks = num_res_blocks

581 self.attn_scales = attn_scales

582 self.temperal_downsample = temperal_downsample

583 self.temperal_upsample = temperal_downsample[::-1]

584

585 # modules

586 self.encoder = Encoder3d(

587 dim, z_dim * 2, dim_mult, num_res_blocks,

588 attn_scales,

589 self.temperal_downsample, dropout)

590 self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)

591 self.conv2 = CausalConv3d(z_dim, z_dim, 1)

592 self.decoder = Decoder3d(

593 dim, z_dim, dim_mult, num_res_blocks,

594 attn_scales,

595 self.temperal_upsample, dropout)

596

597 def forward(

598 self,

599 x: torch.Tensor

600 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

601 mu, log_var = self.encode(x)

602 z = self.reparameterize(mu, log_var)

603 x_recon = self.decode(z)

604 return x_recon, mu, log_var

605

606 def encode_1frame(

607 self,

608 x: torch.Tensor,

609 scale: Tuple[Any, Any]

610 ) -> torch.Tensor:

611 self.clear_cache()

612 # cache

613 t = x.shape[2]

614 iter_ = 1 + (t - 1) // 4

615 # 对encode输入的x，按时间拆分为1、4、4、4....

616 for i in range(iter_):

617 self._enc_conv_idx = [0]

618 if i == 0:

619 out = self.encoder(

620 x[:, :, :1, :, :],

621 feat_cache=self._enc_feat_map,

622 feat_idx=self._enc_conv_idx)

623 else:

624 out_ = self.encoder(

625 x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],

626 feat_cache=self._enc_feat_map,

627 feat_idx=self._enc_conv_idx)

628 out = torch.cat([out, out_], 2)

629 mu, log_var = self.conv1(out).chunk(2, dim=1)

630 if isinstance(scale[0], torch.Tensor):

631 mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(

632 1, self.z_dim, 1, 1, 1)

633 else:

634 mu = (mu - scale[0]) * scale[1]

635 self.clear_cache()

636 return mu

637

638 def decode_1frame(

639 self,

640 z: torch.Tensor,

641 scale: Tuple[Any, Any]

642 ) -> torch.Tensor:

643 self.clear_cache()

644 # z: [b,c,t,h,w]

645 if isinstance(scale[0], torch.Tensor):

646 z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(

647 1, self.z_dim, 1, 1, 1)

648 else:

649 z = z / scale[1] + scale[0]

650 iter_ = z.shape[2]

651 x = self.conv2(z)

652 for i in range(iter_):

653 self._conv_idx = [0]

654 if i == 0:

655 out = self.decoder(

656 x[:, :, i:i + 1, :, :],

657 feat_cache=self._feat_map,

658 feat_idx=self._conv_idx)

659 else:

660 out_ = self.decoder(

661 x[:, :, i:i + 1, :, :],

662 feat_cache=self._feat_map,

663 feat_idx=self._conv_idx)

664 out = torch.cat([out, out_], 2)

665 self.clear_cache()

666 return out

667

668 def encode(

669 self,

670 x: torch.Tensor,

671 scale: Tuple[Any, Any],

672 start_frames: int = 1,

673 end_frames: int = 0

674 ) -> torch.Tensor:

675 self.clear_cache()

676

677 num_frames = x.shape[2]

678 middle_lat_frames = (num_frames - start_frames - end_frames) // 4

679 num_lat_frames = start_frames + middle_lat_frames + end_frames

680

681 # For the x input to the encoder, split it by time into segments of 1, 1,..., 4, 4, 4..., 4, 1,...,1

682 out = []

683 for i in range(num_lat_frames):

684 self._enc_conv_idx = [0]

685 if i < start_frames or i >= num_lat_frames - end_frames:

686 # Unmodified frames at the beginning and end

687 self.clear_cache() # No compression for the start frames

688 frame = x[:, :, i:i + 1, :, :]

689 else:

690 # Compress middle into 4-frame segments

691 i0 = start_frames + 4 * (i - start_frames)

692 i1 = start_frames + 4 * (i - start_frames + 1)

693 frame = x[:, :, i0:i1, :, :]

694

695 frame_encoded = self.encoder(

696 frame,

697 feat_cache=self._enc_feat_map,

698 feat_idx=self._enc_conv_idx)

699 out.append(frame_encoded)

700 out = torch.cat(out, dim=2)

701

702 # Convolution for mu and log_var

703 aux = []

704 # Start frames 1 by 1

705 for i in range(start_frames - 1):

706 out_conv = self.conv1(out[:, :, i:i + 1, :, :])

707 aux.append(out_conv)

708 # Middle frames batched in groups of 4 (after the first one)

709 out_conv = self.conv1(out[:, :, start_frames - 1:num_lat_frames - end_frames, :, :])

710 aux.append(out_conv)

711 # End frames 1 by 1

712 if end_frames > 0:

713 out_conv = self.conv1(out[:, :, -end_frames:, :, :])

714 aux.append(out_conv)

715 mu, log_var = torch.cat(aux, dim=2).chunk(2, dim=1)

716

717 # Scale adjustment

718 if isinstance(scale[0], torch.Tensor):

719 scale0 = scale[0].view(1, self.z_dim, 1, 1, 1)

720 scale1 = scale[1].view(1, self.z_dim, 1, 1, 1)

721 mu = (mu - scale0) * scale1

722 else:

723 mu = (mu - scale[0]) * scale[1]

724

725 self.clear_cache()

726 return mu

727

728 def decode(self, z, scale, start_frames=1, end_frames=0) -> torch.Tensor:

729 self.clear_cache()

730 # z: [b,c,t,h,w]

731 if isinstance(scale[0], torch.Tensor):

732 scale1 = scale[1].view(1, self.z_dim, 1, 1, 1)

733 scale0 = scale[0].view(1, self.z_dim, 1, 1, 1)

734 z = z / scale1 + scale0

735 else:

736 z = z / scale[1] + scale[0]

737 num_lat_frames = z.shape[2]

738

739 # Convolution

740 # x = self.conv2(z) accounting for start/end frames

741 x = []

742 for i in range(0, start_frames - 1):

743 z_conv = self.conv2(z[:, :, i:i + 1, :, :])

744 x.append(z_conv)

745 if end_frames > 0:

746 z_conv = self.conv2(z[:, :, start_frames - 1:-end_frames, :, :])

747 else:

748 z_conv = self.conv2(z[:, :, start_frames - 1:, :, :])

749 x.append(z_conv)

750 if end_frames > 0:

751 for i in range(end_frames):

752 z_conv = self.conv2(z[:, :, -end_frames:-end_frames + 1, :, :])

753 x.append(z_conv)

754 x = torch.cat(x, dim=2)

755

756 # Decode

757 out = []

758 for i in range(num_lat_frames):

759 self._conv_idx = [0]

760 if i < start_frames:

761 self.clear_cache()

762 elif end_frames > 0 and i >= num_lat_frames - end_frames:

763 self.clear_cache()

764

765 frame = x[:, :, i:i + 1, :, :]

766 frame_decoded = self.decoder(

767 frame,

768 feat_cache=self._feat_map,

769 feat_idx=self._conv_idx)

770 out.append(frame_decoded)

771 out = torch.cat(out, 2)

772 self.clear_cache()

773

774 return out

775

776 def decode_stream(

777 self,

778 z: torch.Tensor,

779 scale: Tuple[Any, Any],

780 start_frame: int = 0,

781 start_frames: int = 1,

782 end_frames: int = 0

783 ) -> torch.Tensor:

784 """

785 Decode version that does not wait for all frames and starts yielding them immediately.

786 It yields one frame at a time.

787 z: Latent variable tensor of shape [C, #frames, lat_h, lat_w].

788 """

789 self.clear_cache()

790 # z: [batch, RGB, #lat_frames, lat_h, lat_w]

791 if isinstance(scale[0], torch.Tensor):

792 scale_0 = scale[0].view(1, self.z_dim, 1, 1, 1)

793 scale_1 = scale[1].view(1, self.z_dim, 1, 1, 1)

794 z = z / scale_1 + scale_0

795 else:

796 z = z / scale[1] + scale[0]

797 num_lat_frames = z.shape[2]

798

799 # Convolution

800 # x = self.conv2(z) accounting for start and end frames

801 x = []

802 for i in range(0, start_frames - 1):

803 z_conv = self.conv2(z[:, :, i:i + 1, :, :])

804 x.append(z_conv)

805 if end_frames > 0:

806 z_conv = self.conv2(z[:, :, start_frames - 1:-end_frames, :, :])

807 else:

808 z_conv = self.conv2(z[:, :, start_frames - 1:, :, :])

809 x.append(z_conv)

810 if end_frames > 0:

811 for i in range(end_frames):

812 z_conv = self.conv2(z[:, :, -end_frames:-end_frames + 1, :, :])

813 x.append(z_conv)

814 x = torch.cat(x, dim=2)

815

816 # Decode

817 for i in range(num_lat_frames):

818 self._conv_idx = [0]

819 if i < start_frames:

820 self.clear_cache()

821 elif end_frames > 0 and i >= num_lat_frames - end_frames:

822 self.clear_cache()

823

824 frame = x[:, :, i:i + 1, :, :]

825 out = self.decoder(

826 frame,

827 feat_cache=self._feat_map,

828 feat_idx=self._conv_idx)

829 num_frames = out.shape[2]

830 for cur_frame_ix in range(num_frames):

831 yield out[:, :, cur_frame_ix, :, :] # frames x [batch, RGB, h, w]

832 self.clear_cache()

833

834 def reparameterize(

835 self,

836 mu: torch.Tensor,

837 log_var: torch.Tensor

838 ) -> torch.Tensor:

839 std = torch.exp(0.5 * log_var)

840 eps = torch.randn_like(std)

841 return eps * std + mu

842

843 def sample(

844 self,

845 imgs: torch.Tensor,

846 deterministic: bool = False

847 ) -> torch.Tensor:

848 mu, log_var = self.encode(imgs)

849 if deterministic:

850 return mu

851 std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))

852 return mu + std * torch.randn_like(std)

853

854 def clear_cache(self) -> None:

855 self._conv_num = count_conv3d(self.decoder)

856 self._conv_idx = [0]

857 self._feat_map = [None] * self._conv_num

858 # cache encode

859 self._enc_conv_num = count_conv3d(self.encoder)

860 self._enc_conv_idx = [0]

861 self._enc_feat_map = [None] * self._enc_conv_num

862

863

864def _video_vae(

865 pretrained_path: Optional[str] = None,

866 z_dim: Optional[int] = None,

867 device: str = "cpu",

868 **kwargs: Any

869) -> WanVAE_:

870 """

871 Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.

872 """

873 # params

874 cfg = dict(

875 dim=96,

876 z_dim=z_dim,

877 dim_mult=[1, 2, 4, 4],

878 num_res_blocks=2,

879 attn_scales=[],

880 temperal_downsample=[False, True, True],

881 dropout=0.0)

882 cfg.update(**kwargs)

883

884 # init model

885 with torch.device('meta'):

886 model = WanVAE_(**cfg)

887

888 # load checkpoint

889 logging.info(f'loading {pretrained_path}')

890 model.load_state_dict(

891 torch.load(pretrained_path, map_location=device, weights_only=False), # nosec B614 - trusted model checkpoint

892 assign=True)

894 return model

897class WanVAE:

899 def __init__(

900 self,

901 z_dim: int = 16,

902 vae_pth: str = "cache/vae_step_411000.pth",

903 dtype: torch.dtype = torch.float,

904 device: str = "cuda"

905 ) -> None:

906 self.dtype = dtype

907 self.device = device

908

909 mean = [

910 -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,

911 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921

912 ]

913 std = [

914 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,

915 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160

916 ]

917 self.mean = torch.tensor(mean, dtype=dtype, device=device)

918 self.std = torch.tensor(std, dtype=dtype, device=device)

919 self.scale = [self.mean, 1.0 / self.std]

920

921 # init model

922 self.model = _video_vae(

923 pretrained_path=vae_pth,

924 z_dim=z_dim,

925 ).eval().requires_grad_(False).to(device)

926

927 def encode(

928 self,

929 videos: List[torch.Tensor],

930 start_frames: int = 1,

931 end_frames: int = 0

932 ) -> List[torch.Tensor]:

933 """

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

935 """

936 with amp.autocast('cuda', dtype=self.dtype):

937 return [

938 self.model.encode(u.unsqueeze(0), self.scale, start_frames, end_frames).float().squeeze(0)

939 for u in videos

940 ]

941

942 def decode(

943 self,

944 zs: List[torch.Tensor],

945 start_frames: int = 1,

946 end_frames: int = 0

947 ) -> List[torch.Tensor]:

948 with amp.autocast('cuda', dtype=self.dtype):

949 return [

950 self.model.decode(

951 u.unsqueeze(0),

952 self.scale,

953 start_frames,

954 end_frames

955 ).float().clamp_(-1, 1).squeeze(0)

956 for u in zs

957 ]

958

959 def decode_stream(

960 self,

961 z: torch.Tensor,

962 start_frame: int = 0,

963 start_frames: int = 1,

964 end_frames: int = 0

965 ) -> Generator[torch.Tensor, None, None]:

966 """

967 Decode a video in latent space into a regular video yielding frame by frame.

968 z: Latent variable tensor of shape [C, #frames, lat_h, lat_w] (e.g., [16, 1+20, 68, 90]).

969 num_frames: Number of frames to decode.

970 """

971 with amp.autocast('cuda', dtype=self.dtype):

972 for video_frame in self.model.decode_stream(

973 z.unsqueeze(0),

974 self.scale,

975 start_frame=start_frame,

976 start_frames=start_frames,

977 end_frames=end_frames

978 ):

979 yield video_frame.float().clamp_(-1, 1).squeeze(0)