def compute_density_for_timestep_sampling_scm(batch_size: int, logit_mean: float = None, logit_std: float = None):
    """Compute the density for sampling the timesteps when doing Sana-Sprint training."""
    sigma = torch.randn(batch_size, device="cpu")
    sigma = (sigma * logit_std + logit_mean).exp()
    u = torch.atan(sigma / 0.5)  # TODO: 0.5 should be a hyper-parameter

    return u
Compute the density for sampling the timesteps when doing Sana-Sprint training.
examples/research_projects/sana/train_sana_sprint_diffusers.py