Function modulate

models/lightningdit.py:103–106

(x, shift, scale)

    return x * (1 + scale) + shift
@torch.compile
def modulate(x, shift, scale):
    if shift is None:
        return x * (1 + scale.unsqueeze(1))
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class Attention(nn.Module):
    """