infworld/models/scheduler.py

import math
import time
import numpy as np

from tqdm import tqdm
from typing import Callable
from einops import rearrange
from functools import partial

import torch
from torch.distributions import LogisticNormal

from infworld.context_parallel import context_parallel_util

# some code are inspired by https://github.com/magic-research/piecewise-rectified-flow/blob/main/scripts/train_perflow.py
# and https://github.com/magic-research/piecewise-rectified-flow/blob/main/src/scheduler_perflow.py
# and https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py


def _extract_into_tensor(arr, timesteps, broadcast_shape):
    """
    Extract values from a 1-D numpy array for a batch of indices.
    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
    while len(res.shape) < len(broadcast_shape):
        res = res[..., None]
    return res + torch.zeros(broadcast_shape, device=timesteps.device)


def mean_flat(tensor: torch.Tensor, stoploss_mask=None):
    """
    Take the mean over all non-batch dimensions.
        tensor: [B, C, T, H, W]
        stoploss_mask: [B, T, H, W]
    """
    if stoploss_mask is None:
        return tensor.mean(dim=list(range(1, len(tensor.shape))))
    else:
        stoploss_mask = stoploss_mask.unsqueeze(1).expand_as(tensor) # [B, T, H, W] --> [B, C, T, H, W]
        assert tensor.shape == stoploss_mask.shape, f"shape of tensor {tensor.shape} and stoploss_mask {stoploss_mask.shape} should be the same"
        loss_mask = ~stoploss_mask
        masked_loss = tensor * loss_mask
        sum_loss = masked_loss.sum(dim=list(range(1, len(tensor.shape))))
        count_nonzero = loss_mask.sum(dim=list(range(1, len(tensor.shape))))
        mean_loss = sum_loss / count_nonzero.clamp(min=1)

        return mean_loss


def clamp(value, min_value, max_value):
    return max(min_value, min(value, max_value))



def timestep_transform(
    t,
    shift=5.0,
    num_timesteps=1000,
):
    t = t / num_timesteps
    # shift the timestep based on ratio
    new_t = shift * t / (1 + (shift - 1) * t)
    new_t = new_t * num_timesteps
    return new_t


class RFlowScheduler:
    def __init__(
        self,
        num_timesteps=1000,
        num_sampling_steps=10,
        use_discrete_timesteps=False,
        sample_method="uniform",
        loc=0.0,
        scale=1.0,
        shift=5.0,
        use_timestep_transform=False,
        transform_scale=1.0,
        use_reversed_velocity=False,
        cfg_scale=7.0,
        **kwargs,
    ):
        self.num_timesteps = num_timesteps
        self.num_sampling_steps = num_sampling_steps
        self.use_discrete_timesteps = use_discrete_timesteps
        self.use_reversed_velocity = use_reversed_velocity
        self.cfg_scale = cfg_scale

        # sample method
        assert sample_method in ["uniform", "logit-normal"]
        assert (
            sample_method == "uniform" or not use_discrete_timesteps
        ), "Only uniform sampling is supported for discrete timesteps"
        self.sample_method = sample_method
        if sample_method == "logit-normal":
            self.distribution = LogisticNormal(torch.tensor([loc]), torch.tensor([scale]))
            self.sample_t = lambda x: self.distribution.sample((x.shape[0],))[:, 0].to(x.device)

        # timestep transform
        self.use_timestep_transform = use_timestep_transform
        self.transform_scale = transform_scale

        self.shift = shift
        sigmas = torch.linspace(0, 1, num_timesteps)
        sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
        self.timesteps = sigmas * num_timesteps

        y = torch.exp(-2 * ((self.timesteps - num_timesteps/2) / num_timesteps)**2)
        y_shifted = y - y.min()
        self.bsmntw_weighing = y_shifted * (num_timesteps / y_shifted.sum())

    def training_losses(self, model, x_start, model_kwargs=None, noise=None, x_ignore_mask=None, t=None):
        """
        Compute training losses for a single timestep.
        Arguments format copied from opensora/schedulers/iddpm/gaussian_diffusion.py/training_losses
        Note: t is int tensor and should be rescaled from [0, num_timesteps-1] to [1,0]
        """
        
        if t is None:
            if self.use_discrete_timesteps:
                t = torch.randint(0, self.num_timesteps, (x_start.shape[0],), device=x_start.device)
            elif self.sample_method == "uniform":
                t = torch.rand((x_start.shape[0],), device=x_start.device) * self.num_timesteps
            elif self.sample_method == "logit-normal":
                t = self.sample_t(x_start) * self.num_timesteps

            if self.use_timestep_transform:
                latent_size = x_start.shape[-3:]
                t = timestep_transform(t, shift=self.shift, num_timesteps=self.num_timesteps)

        if model_kwargs is None:
            model_kwargs = {}
        if noise is None:
            noise = torch.randn_like(x_start)
        assert noise.shape == x_start.shape


        if context_parallel_util.get_cp_size() > 1:
            context_parallel_util.cp_broadcast(noise)
            context_parallel_util.cp_broadcast(t)

        x_t = self.add_noise(x_start, noise, t)


        target = x_start - noise
        if self.use_reversed_velocity:
            target = -target

        terms = {}
        model_output = model(x_t, t, x_ignore_mask=x_ignore_mask, **model_kwargs)
        velocity_pred = model_output

        T = target.shape[2]
        loss = mean_flat((velocity_pred[:, :, -T:] - target).pow(2), stoploss_mask=x_ignore_mask[:, -T:])

        # # get loss weight
        # timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - t.unsqueeze(1).to(self.timesteps.device)).abs(), dim=1)
        # weights = self.bsmntw_weighing[timestep_id]
        # loss = weights.to(loss) * loss

        terms["loss"] = loss

        return terms

    def add_noise(
        self,
        original_samples: torch.FloatTensor,
        noise: torch.FloatTensor,
        timesteps: torch.IntTensor,
    ) -> torch.FloatTensor:
        """
        compatible with diffusers add_noise()
        """
        timesteps = timesteps.float() / self.num_timesteps
        timesteps = timesteps.view(timesteps.shape + (1,) * (len(noise.shape)-1))

        return (1 - timesteps) * original_samples + timesteps * noise

    def sample(
        self,
        model,
        text_encoder,
        null_embedder,
        z_size,
        prompts,
        device,
        mask=None,
        guidance_scale=None,
        negative_prompts=None,
        additional_args=None,
        progress=True,
    ):
        # if no specific guidance scale is provided, use the default scale when initializing the scheduler
        if guidance_scale is None:
            guidance_scale = self.cfg_scale

        n = len(prompts)
        z = torch.randn(*z_size, device=device)

        if context_parallel_util.get_cp_size() > 1:
            context_parallel_util.cp_broadcast(z)
        
        # For performance alignment
        # from source.opensora.utils.inference_utils import apply_mask_strategy
        # mask = apply_mask_strategy(z, [[]], [""], 0, align=5)
        
        assert negative_prompts is None or len(negative_prompts) in [n, 1], \
            "Invalid negative prompts."

        if negative_prompts:
            if len(negative_prompts) == 1: negative_prompts *= n
            prompts = prompts + negative_prompts

        batch_size = len(prompts)
        if context_parallel_util.get_cp_rank() == 0:
            model_args = text_encoder.encode(prompts)
            if context_parallel_util.get_cp_size() > 1:
                context_parallel_util.cp_broadcast(model_args['y'])
                context_parallel_util.cp_broadcast(model_args['y_mask'])
        elif context_parallel_util.get_cp_size() > 1:
            caption_channels = text_encoder.output_dim
            model_max_length = text_encoder.model_max_length
            y_tensor = torch.zeros([batch_size, 1, model_max_length, caption_channels], dtype=torch.float32, device=device)
            y_mask_tensor = torch.zeros([batch_size, model_max_length], dtype=torch.int64, device=device)
            context_parallel_util.cp_broadcast(y_tensor)
            context_parallel_util.cp_broadcast(y_mask_tensor)
            model_args = {
                "y" : y_tensor,
                "y_mask": y_mask_tensor,
            }

        assert negative_prompts, "Not support uncond training now, pls use negative prompt for uncond."
        if not negative_prompts:
            uncond = null_embedder.y_embedding[None].repeat(n, 1, 1)[:, None]
            model_args["y"] = torch.concat([model_args["y"], uncond])

        if additional_args is not None:
            model_args.update(additional_args)

        # prepare timesteps
        timesteps = list(np.linspace(self.num_timesteps, 1, self.num_sampling_steps, dtype=np.float32))
        if self.use_discrete_timesteps:
            timesteps = [int(round(t)) for t in timesteps]
        timesteps = [torch.tensor([t] * z.shape[0], device=device) for t in timesteps]
        if self.use_timestep_transform:
            latent_size = z_size[-3:]
            timesteps = [timestep_transform(t, shift=self.shift, num_timesteps=self.num_timesteps) for t in timesteps]

        if mask is not None:
            noise_added = torch.zeros_like(mask, dtype=torch.bool)
            noise_added = noise_added | (mask == 1)

        if context_parallel_util.get_cp_size() > 1:
            torch.distributed.barrier(group=context_parallel_util.get_cp_group())

        model_args["image_cond"] = model_args["image_cond"].repeat(2, 1, 1, 1, 1)
        progress_wrap = partial(tqdm, total=len(timesteps)) if progress else (lambda x: x)
        for i, t in progress_wrap(enumerate(timesteps)):
            # mask for adding noise
            if mask is not None:
                mask_t = mask * self.num_timesteps
                x0 = z.clone()

                x0_noise = torch.randn_like(x0)
                if context_parallel_util.get_cp_size() > 1:
                    context_parallel_util.cp_broadcast(x0_noise)

                x_noise = self.scheduler.add_noise(x0, x0_noise, t)

                mask_t_upper = mask_t >= t.unsqueeze(1)
                model_args["x_mask"] = mask_t_upper.repeat(2, 1)
                mask_add_noise = mask_t_upper & ~noise_added

                z = torch.where(mask_add_noise[:, None, :, None, None], x_noise, x0)
                noise_added = mask_t_upper

            # classifier-free guidance
            z_in = torch.cat([z, z], 0)

            t = torch.cat([t, t], 0)
            start = time.time()
            pred = model(z_in, t, **model_args)
            pred = pred[:, :, -z_in.shape[2]:]
            end = time.time()

            print(f"Step {i} Forward time: {end - start:.4f} seconds")
            pred_cond, pred_uncond = pred.chunk(2, dim=0)
            v_pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond)

            # When model predict noise-z0, the actual velocity is (v_pred * -1)
            if self.use_reversed_velocity:
                v_pred = -v_pred

            # update z
            dt = timesteps[i] - timesteps[i + 1] if i < len(timesteps) - 1 else timesteps[i]
            dt = dt / self.num_timesteps
            z = z + v_pred * dt[:, None, None, None, None]

            if mask is not None:
                z = torch.where(mask_t_upper[:, None, :, None, None], z, x0)
        return z