Method forward

k_diffusion/models/modules.py:587–592

(self, x)

        self.out_norm = RMSNorm(d_model)

    def forward(self, x):
        x = self.in_norm(x)
        for block in self.blocks:
            x = block(x)
        x = self.out_norm(x)
        return x
    

# Token merging and splitting