class TimestepBlock(nn.Module):
    def __init__(self, c, c_timestep):
        super().__init__()

        self.mapper = nn.Linear(c_timestep, c * 2)

    def forward(self, x, t):
        a, b = self.mapper(t)[:, :, None, None].chunk(2, dim=1)
        return x * (1 + a) + b
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py