36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import types
|
|
|
|
import torch
|
|
|
|
from .forwards import (
|
|
block_forward,
|
|
decoder_forward,
|
|
res_block_forward,
|
|
upsample_forward,
|
|
)
|
|
|
|
|
|
def patch_vae(vae_model, patch_block=4):
|
|
|
|
vae_model.decoder.forward = types.MethodType(decoder_forward, vae_model.decoder)
|
|
vae_model.decoder.conv_out.conv.weight.data = (
|
|
vae_model.decoder.conv_out.conv.weight.data.to(
|
|
memory_format=torch.channels_last_3d
|
|
)
|
|
)
|
|
|
|
for i, block in enumerate(reversed(vae_model.decoder.up_blocks)):
|
|
if i >= patch_block:
|
|
break
|
|
if block.__class__.__name__ == "UNetMidBlock3D":
|
|
block.forward = types.MethodType(block_forward, block)
|
|
for name, param in block.named_parameters():
|
|
if "conv1" in name or "conv2" in name:
|
|
if "weight" in name:
|
|
param = param.to(memory_format=torch.channels_last_3d)
|
|
for res_block in block.res_blocks:
|
|
if res_block.__class__.__name__ == "ResnetBlock3D":
|
|
res_block.forward = types.MethodType(res_block_forward, res_block)
|
|
elif block.__class__.__name__ == "DepthToSpaceUpsample":
|
|
block.forward = types.MethodType(upsample_forward, block)
|