Method __init__

models/unet_model_google.py:180–202

(
        self,
        dim,
        dim_out,
        *,
        time_emb_dim=None,
        norm_groups=32,
        dropout=0,
        with_attn=False,
        attn_guide=False,
    )

class ResnetBlocWithAttn(nn.Module):
    def __init__(
        self,
        dim,
        dim_out,
        *,
        time_emb_dim=None,
        norm_groups=32,
        dropout=0,
        with_attn=False,
        attn_guide=False,
    ):
        super().__init__()
        self.with_attn = with_attn
        self.res_block = ResnetBlock(
            dim,
            dim_out,
            time_emb_dim,
            norm_groups=norm_groups,
            dropout=dropout,
            attn_guide=attn_guide,
        )
        if with_attn:
            self.attn = SelfAttention(dim_out, norm_groups=norm_groups, nhead=8)

    def forward(self, x, time_emb, guidance=None):
        x = self.res_block(x, time_emb, guidance)