(self)
| 7 | |
| 8 | class SD3VAEEncoder(torch.nn.Module): |
| 9 | def __init__(self): |
| 10 | super().__init__() |
| 11 | self.scaling_factor = 1.5305 # Different from SD 1.x |
| 12 | self.shift_factor = 0.0609 # Different from SD 1.x |
| 13 | self.conv_in = torch.nn.Conv2d(3, 128, kernel_size=3, padding=1) |
| 14 | |
| 15 | self.blocks = torch.nn.ModuleList([ |
| 16 | # DownEncoderBlock2D |
| 17 | ResnetBlock(128, 128, eps=1e-6), |
| 18 | ResnetBlock(128, 128, eps=1e-6), |
| 19 | DownSampler(128, padding=0, extra_padding=True), |
| 20 | # DownEncoderBlock2D |
| 21 | ResnetBlock(128, 256, eps=1e-6), |
| 22 | ResnetBlock(256, 256, eps=1e-6), |
| 23 | DownSampler(256, padding=0, extra_padding=True), |
| 24 | # DownEncoderBlock2D |
| 25 | ResnetBlock(256, 512, eps=1e-6), |
| 26 | ResnetBlock(512, 512, eps=1e-6), |
| 27 | DownSampler(512, padding=0, extra_padding=True), |
| 28 | # DownEncoderBlock2D |
| 29 | ResnetBlock(512, 512, eps=1e-6), |
| 30 | ResnetBlock(512, 512, eps=1e-6), |
| 31 | # UNetMidBlock2D |
| 32 | ResnetBlock(512, 512, eps=1e-6), |
| 33 | VAEAttentionBlock(1, 512, 512, 1, eps=1e-6), |
| 34 | ResnetBlock(512, 512, eps=1e-6), |
| 35 | ]) |
| 36 | |
| 37 | self.conv_norm_out = torch.nn.GroupNorm(num_channels=512, num_groups=32, eps=1e-6) |
| 38 | self.conv_act = torch.nn.SiLU() |
| 39 | self.conv_out = torch.nn.Conv2d(512, 32, kernel_size=3, padding=1) |
| 40 | |
| 41 | def tiled_forward(self, sample, tile_size=64, tile_stride=32): |
| 42 | hidden_states = TileWorker().tiled_forward( |
nothing calls this directly
no test coverage detected