Method __init__

deepfloyd_if/model/unet.py:670–679

(self, low_res_diffusion, interpolate_mode='bilinear', *args, **kwargs)

    """

    def __init__(self, low_res_diffusion, interpolate_mode='bilinear', *args, **kwargs):
        self.low_res_diffusion = low_res_diffusion
        self.interpolate_mode = interpolate_mode
        super().__init__(*args, **kwargs)

        self.aug_proj = nn.Sequential(
            linear(self.model_channels, self.time_embed_dim, dtype=self.dtype),
            get_activation(kwargs['activation']),
            linear(self.time_embed_dim, self.time_embed_dim, dtype=self.dtype),
        )

    def forward(self, x, timesteps, low_res, aug_level=None, **kwargs):
        bs, _, new_height, new_width = x.shape