Class MidBlock

lwm/vqgan.py:340–351

class MidBlock(nn.Module):
    config: VQGANConfig
    no_attn: bool
    dropout: float

    @nn.compact
    def __call__(self, hidden_states):
        hidden_states = ResnetBlock(dropout_prob=self.dropout)(hidden_states)
        if not self.no_attn:
            hidden_states = AttnBlock()(hidden_states)
        hidden_states = ResnetBlock(dropout_prob=self.dropout)(hidden_states)
        return hidden_states