models/samplers/ddpm.py

import torch


def ddpm_sampler(
    net,
    batch,
    conditioning_keys=None,
    scheduler=None,
    uncond_tokens=None,
    num_steps=1000,
    cfg_rate=0,
    generator=None,
    use_confidence_sampling=False,
    use_uncond_token=True,
    confidence_value=1.0,
    unconfidence_value=0.0,
):
    if scheduler is None:
        raise ValueError("Scheduler must be provided")

    x_cur = batch["y"].to(torch.float32)
    latents = batch["previous_latents"]
    if use_confidence_sampling:
        batch["confidence"] = (
            torch.ones(x_cur.shape[0], device=x_cur.device) * confidence_value
        )
    step_indices = torch.arange(num_steps + 1, dtype=torch.float32, device=x_cur.device)
    steps = 1 - step_indices / num_steps
    gammas = scheduler(steps)
    latents_cond = latents_uncond = latents
    # dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    dtype = torch.float32
    if cfg_rate > 0 and conditioning_keys is not None:
        stacked_batch = {}
        for key in conditioning_keys:
            if f"{key}_mask" in batch:
                if use_confidence_sampling and not use_uncond_token:
                    stacked_batch[f"{key}_mask"] = torch.cat(
                        [batch[f"{key}_mask"], batch[f"{key}_mask"]], dim=0
                    )
                else:
                    if (
                        batch[f"{key}_mask"].shape[1]
                        > uncond_tokens[f"{key}_mask"].shape[1]
                    ):
                        uncond_mask = (
                            torch.zeros_like(batch[f"{key}_mask"])
                            if batch[f"{key}_mask"].dtype == torch.bool
                            else torch.ones_like(batch[f"{key}_mask"]) * -torch.inf
                        )
                        uncond_mask[:, : uncond_tokens[f"{key}_mask"].shape[1]] = (
                            uncond_tokens[f"{key}_mask"]
                        )
                    else:
                        uncond_mask = uncond_tokens[f"{key}_mask"]
                        batch[f"{key}_mask"] = torch.cat(
                            [
                                batch[f"{key}_mask"],
                                torch.zeros(
                                    batch[f"{key}_mask"].shape[0],
                                    uncond_tokens[f"{key}_embeddings"].shape[1]
                                    - batch[f"{key}_mask"].shape[1],
                                    device=batch[f"{key}_mask"].device,
                                    dtype=batch[f"{key}_mask"].dtype,
                                ),
                            ],
                            dim=1,
                        )
                    stacked_batch[f"{key}_mask"] = torch.cat(
                        [batch[f"{key}_mask"], uncond_mask], dim=0
                    )
            if f"{key}_embeddings" in batch:
                if use_confidence_sampling and not use_uncond_token:
                    stacked_batch[f"{key}_embeddings"] = torch.cat(
                        [
                            batch[f"{key}_embeddings"],
                            batch[f"{key}_embeddings"],
                        ],
                        dim=0,
                    )
                else:
                    if (
                        batch[f"{key}_embeddings"].shape[1]
                        > uncond_tokens[f"{key}_embeddings"].shape[1]
                    ):
                        uncond_tokens[f"{key}_embeddings"] = torch.cat(
                            [
                                uncond_tokens[f"{key}_embeddings"],
                                torch.zeros(
                                    uncond_tokens[f"{key}_embeddings"].shape[0],
                                    batch[f"{key}_embeddings"].shape[1]
                                    - uncond_tokens[f"{key}_embeddings"].shape[1],
                                    uncond_tokens[f"{key}_embeddings"].shape[2],
                                    device=uncond_tokens[f"{key}_embeddings"].device,
                                ),
                            ],
                            dim=1,
                        )
                    elif (
                        batch[f"{key}_embeddings"].shape[1]
                        < uncond_tokens[f"{key}_embeddings"].shape[1]
                    ):
                        batch[f"{key}_embeddings"] = torch.cat(
                            [
                                batch[f"{key}_embeddings"],
                                torch.zeros(
                                    batch[f"{key}_embeddings"].shape[0],
                                    uncond_tokens[f"{key}_embeddings"].shape[1]
                                    - batch[f"{key}_embeddings"].shape[1],
                                    batch[f"{key}_embeddings"].shape[2],
                                    device=batch[f"{key}_embeddings"].device,
                                ),
                            ],
                            dim=1,
                        )
                    stacked_batch[f"{key}_embeddings"] = torch.cat(
                        [
                            batch[f"{key}_embeddings"],
                            uncond_tokens[f"{key}_embeddings"],
                        ],
                        dim=0,
                    )
            elif key not in batch:
                raise ValueError(f"Key {key} not in batch")
            else:
                if isinstance(batch[key], torch.Tensor):
                    if use_confidence_sampling and not use_uncond_token:
                        stacked_batch[key] = torch.cat([batch[key], batch[key]], dim=0)
                    else:
                        stacked_batch[key] = torch.cat(
                            [batch[key], uncond_tokens], dim=0
                        )
                elif isinstance(batch[key], list):
                    if use_confidence_sampling and not use_uncond_token:
                        stacked_batch[key] = [*batch[key], *batch[key]]
                    else:
                        stacked_batch[key] = [*batch[key], *uncond_tokens]
                else:
                    raise ValueError(
                        "Conditioning must be a tensor or a list of tensors"
                    )
        if use_confidence_sampling:
            stacked_batch["confidence"] = torch.cat(
                [
                    torch.ones(x_cur.shape[0], device=x_cur.device) * confidence_value,
                    torch.ones(x_cur.shape[0], device=x_cur.device)
                    * unconfidence_value,
                ],
                dim=0,
            )
    for step, (gamma_now, gamma_next) in enumerate(zip(gammas[:-1], gammas[1:])):
        with torch.cuda.amp.autocast(dtype=dtype):
            if cfg_rate > 0 and conditioning_keys is not None:
                stacked_batch["y"] = torch.cat([x_cur, x_cur], dim=0)
                stacked_batch["gamma"] = gamma_now.expand(x_cur.shape[0] * 2)
                stacked_batch["previous_latents"] = (
                    torch.cat([latents_cond, latents_uncond], dim=0)
                    if latents is not None
                    else None
                )
                denoised_all, latents_all = net(stacked_batch)
                denoised_cond, denoised_uncond = denoised_all.chunk(2, dim=0)
                latents_cond, latents_uncond = latents_all.chunk(2, dim=0)
                denoised = denoised_cond * (1 + cfg_rate) - denoised_uncond * cfg_rate
            else:
                batch["y"] = x_cur
                batch["gamma"] = gamma_now.expand(x_cur.shape[0])
                batch["previous_latents"] = latents
                denoised, latents = net(
                    batch,
                )
        x_pred = (x_cur - torch.sqrt(1 - gamma_now) * denoised) / torch.sqrt(gamma_now)
        x_pred = torch.clamp(x_pred, -1, 1)
        noise_pred = (x_cur - torch.sqrt(gamma_now) * x_pred) / torch.sqrt(
            1 - gamma_now
        )

        log_alpha_t = torch.log(gamma_now) - torch.log(gamma_next)
        alpha_t = torch.clip(torch.exp(log_alpha_t), 0, 1)
        x_mean = torch.rsqrt(alpha_t) * (
            x_cur - torch.rsqrt(1 - gamma_now) * (1 - alpha_t) * noise_pred
        )
        var_t = 1 - alpha_t
        eps = torch.randn(x_cur.shape, device=x_cur.device, generator=generator)
        x_next = x_mean + torch.sqrt(var_t) * eps
        x_cur = x_next
    return x_cur.to(torch.float32)