Function modulate

k_diffusion/models/modules.py:38–39

(x, shift, scale)

# Helpers

def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1).unsqueeze(1)) + shift.unsqueeze(1).unsqueeze(1)

def zero_init(layer):
    nn.init.zeros_(layer.weight)