Class Timestep

sat/sgm/modules/diffusionmodules/openaimodel.py:485–491

class Timestep(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        return timestep_embedding(t, self.dim)


str_to_dtype = {"fp32": th.float32, "fp16": th.float16, "bf16": th.bfloat16}