Open-Sora-v2

Running on Zero

Open-Sora-v2 / opensora /utils /sampling.py

Upload folder using huggingface_hub

3d736fe verified about 2 months ago

22.7 kB

import math
import os
import random
from abc import ABC, abstractmethod
from dataclasses import dataclass, replace

import torch
from einops import rearrange, repeat
from mmengine.config import Config
from peft import PeftModel
from torch import Tensor, nn

from opensora.datasets.aspect import get_image_size
from opensora.models.mmdit.model import MMDiTModel
from opensora.models.text.conditioner import HFEmbedder
from opensora.registry import MODELS, build_module
from opensora.utils.inference import (
    SamplingMethod,
    collect_references_batch,
    prepare_inference_condition,
)

# ======================================================
# Sampling Options
# ======================================================


@dataclass
class SamplingOption:
    # The width of the image/video.
    width: int | None = None

    # The height of the image/video.
    height: int | None = None

    # The resolution of the image/video. If provided, it will override the height and width.
    resolution: str | None = None

    # The aspect ratio of the image/video. If provided, it will override the height and width.
    aspect_ratio: str | None = None

    # The number of frames.
    num_frames: int = 1

    # The number of sampling steps.
    num_steps: int = 50

    # The classifier-free guidance (text).
    guidance: float = 4.0

    # use oscillation for text guidance
    text_osci: bool = False

    # The classifier-free guidance (image), or for the guidance on condition for i2v and v2v
    guidance_img: float | None = None

    # use oscillation for image guidance
    image_osci: bool = False

    # use temporal scaling for image guidance
    scale_temporal_osci: bool = False

    # The seed for the random number generator.
    seed: int | None = None

    # Whether to shift the schedule.
    shift: bool = True

    # The sampling method.
    method: str | SamplingMethod = SamplingMethod.I2V

    # Temporal reduction
    temporal_reduction: int = 1

    # is causal vae
    is_causal_vae: bool = False

    # flow shift
    flow_shift: float | None = None


def sanitize_sampling_option(sampling_option: SamplingOption) -> SamplingOption:
    """
    Sanitize the sampling options.
    Args:
        sampling_option (SamplingOption): The sampling options.
    Returns:
        SamplingOption: The sanitized sampling options.
    """
    if (
        sampling_option.resolution is not None
        or sampling_option.aspect_ratio is not None
    ):
        assert (
            sampling_option.resolution is not None
            and sampling_option.aspect_ratio is not None
        ), "Both resolution and aspect ratio must be provided"
        resolution = sampling_option.resolution
        aspect_ratio = sampling_option.aspect_ratio
        height, width = get_image_size(resolution, aspect_ratio, training=False)
    else:
        assert (
            sampling_option.height is not None and sampling_option.width is not None
        ), "Both height and width must be provided"
        height, width = sampling_option.height, sampling_option.width

    height = (height // 16 + (1 if height % 16 else 0)) * 16
    width = (width // 16 + (1 if width % 16 else 0)) * 16
    replace_dict = dict(height=height, width=width)

    if isinstance(sampling_option.method, str):
        method = SamplingMethod(sampling_option.method)
        replace_dict["method"] = method

    return replace(sampling_option, **replace_dict)


def get_oscillation_gs(guidance_scale: float, i: int, force_num=10):
    """
    get oscillation guidance for cfg.
    Args:
        guidance_scale: original guidance value
        i: denoising step
        force_num: before which don't apply oscillation
    """
    if i < force_num or (i >= force_num and i % 2 == 0):
        gs = guidance_scale
    else:
        gs = 1.0
    return gs


# ======================================================
# Denoising
# ======================================================


class Denoiser(ABC):
    @abstractmethod
    def denoise(self, model: MMDiTModel, **kwargs) -> Tensor:
        """Denoise the input."""

    @abstractmethod
    def prepare_guidance(
        self,
        text: list[str],
        optional_models: dict[str, nn.Module],
        device: torch.device,
        dtype: torch.dtype,
        **kwargs,
    ) -> dict[str, Tensor]:
        """Prepare the guidance for the model. This method will alter text."""


class I2VDenoiser(Denoiser):
    def denoise(self, model: MMDiTModel, **kwargs) -> Tensor:
        img = kwargs.pop("img")
        timesteps = kwargs.pop("timesteps")
        guidance = kwargs.pop("guidance")
        guidance_img = kwargs.pop("guidance_img")

        # cond ref arguments
        masks = kwargs.pop("masks")
        masked_ref = kwargs.pop("masked_ref")
        kwargs.pop("sigma_min")

        # oscillation guidance
        text_osci = kwargs.pop("text_osci", False)
        image_osci = kwargs.pop("image_osci", False)
        scale_temporal_osci = kwargs.pop("scale_temporal_osci", False)

        # patch size
        patch_size = kwargs.pop("patch_size", 2)

        guidance_vec = torch.full(
            (img.shape[0],), guidance, device=img.device, dtype=img.dtype
        )
        for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
            # timesteps
            t_vec = torch.full(
                (img.shape[0],), t_curr, dtype=img.dtype, device=img.device
            )
            b, c, t, w, h = masked_ref.size()
            cond = torch.cat((masks, masked_ref), dim=1)
            cond = pack(cond, patch_size=patch_size)
            kwargs["cond"] = torch.cat([cond, cond, torch.zeros_like(cond)], dim=0)

            # forward preparation
            cond_x = img[: len(img) // 3]

            img = torch.cat([cond_x, cond_x, cond_x], dim=0)
            # forward
            pred = model(
                img=img,
                **kwargs,
                timesteps=t_vec,
                guidance=guidance_vec,
            )

            # prepare guidance
            text_gs = get_oscillation_gs(guidance, i) if text_osci else guidance
            image_gs = (
                get_oscillation_gs(guidance_img, i) if image_osci else guidance_img
            )
            cond, uncond, uncond_2 = pred.chunk(3, dim=0)
            if image_gs > 1.0 and scale_temporal_osci:
                # image_gs decrease with each denoising step
                step_upper_image_gs = torch.linspace(image_gs, 1.0, len(timesteps))[i]
                # image_gs increase along the temporal axis of the latent video
                image_gs = torch.linspace(1.0, step_upper_image_gs, t)[
                    None, None, :, None, None
                ].repeat(b, c, 1, h, w)
                image_gs = pack(image_gs, patch_size=patch_size).to(cond.device, cond.dtype)

            # update
            pred = uncond_2 + image_gs * (uncond - uncond_2) + text_gs * (cond - uncond)
            pred = torch.cat([pred, pred, pred], dim=0)

            img = img + (t_prev - t_curr) * pred

        img = img[: len(img) // 3]

        return img

    def prepare_guidance(
        self,
        text: list[str],
        optional_models: dict[str, nn.Module],
        device: torch.device,
        dtype: torch.dtype,
        **kwargs,
    ) -> tuple[list[str], dict[str, Tensor]]:
        ret = {}

        neg = kwargs.get("neg", None)
        ret["guidance_img"] = kwargs.pop("guidance_img")

        # text
        if neg is None:
            neg = [""] * len(text)
        text = text + neg + neg
        return text, ret


class DistilledDenoiser(Denoiser):
    def denoise(self, model: MMDiTModel, **kwargs) -> Tensor:
        img = kwargs.pop("img")
        timesteps = kwargs.pop("timesteps")
        guidance = kwargs.pop("guidance")

        guidance_vec = torch.full(
            (img.shape[0],), guidance, device=img.device, dtype=img.dtype
        )
        for t_curr, t_prev in zip(timesteps[:-1], timesteps[1:]):
            # timesteps
            t_vec = torch.full(
                (img.shape[0],), t_curr, dtype=img.dtype, device=img.device
            )
            # forward
            pred = model(
                img=img,
                **kwargs,
                timesteps=t_vec,
                guidance=guidance_vec,
            )
            # update
            img = img + (t_prev - t_curr) * pred
        return img

    def prepare_guidance(
        self,
        text: list[str],
        optional_models: dict[str, nn.Module],
        device: torch.device,
        dtype: torch.dtype,
        **kwargs,
    ) -> tuple[list[str], dict[str, Tensor]]:
        return text, {}


SamplingMethodDict = {
    SamplingMethod.I2V: I2VDenoiser(),
    SamplingMethod.DISTILLED: DistilledDenoiser(),
}


# ======================================================
# Timesteps
# ======================================================


def time_shift(alpha: float, t: Tensor) -> Tensor:
    return alpha * t / (1 + (alpha - 1) * t)


def get_res_lin_function(
    x1: float = 256, y1: float = 1, x2: float = 4096, y2: float = 3
) -> callable:
    m = (y2 - y1) / (x2 - x1)
    b = y1 - m * x1
    return lambda x: m * x + b


def get_schedule(
    num_steps: int,
    image_seq_len: int,
    num_frames: int,
    shift_alpha: float | None = None,
    base_shift: float = 1,
    max_shift: float = 3,
    shift: bool = True,
) -> list[float]:
    # extra step for zero
    timesteps = torch.linspace(1, 0, num_steps + 1)

    # shifting the schedule to favor high timesteps for higher signal images
    if shift:
        if shift_alpha is None:
            # estimate mu based on linear estimation between two points
            # spatial scale
            shift_alpha = get_res_lin_function(y1=base_shift, y2=max_shift)(
                image_seq_len
            )
            # temporal scale
            shift_alpha *= math.sqrt(num_frames)
        # calculate shifted timesteps
        timesteps = time_shift(shift_alpha, timesteps)

    return timesteps.tolist()


def get_noise(
    num_samples: int,
    height: int,
    width: int,
    num_frames: int,
    device: torch.device,
    dtype: torch.dtype,
    seed: int,
    patch_size: int = 2,
    channel: int = 16,
) -> Tensor:
    """
    Generate a noise tensor.
    Args:
        num_samples (int): Number of samples.
        height (int): Height of the noise tensor.
        width (int): Width of the noise tensor.
        num_frames (int): Number of frames.
        device (torch.device): Device to put the noise tensor on.
        dtype (torch.dtype): Data type of the noise tensor.
        seed (int): Seed for the random number generator.
    Returns:
        Tensor: The noise tensor.
    """
    D = int(os.environ.get("AE_SPATIAL_COMPRESSION", 16))
    return torch.randn(
        num_samples,
        channel,
        num_frames,
        # allow for packing
        patch_size * math.ceil(height / D),
        patch_size * math.ceil(width / D),
        device=device,
        dtype=dtype,
        generator=torch.Generator(device=device).manual_seed(seed),
    )


def pack(x: Tensor, patch_size: int = 2) -> Tensor:
    return rearrange(
        x, "b c t (h ph) (w pw) -> b (t h w) (c ph pw)", ph=patch_size, pw=patch_size
    )


def unpack(
    x: Tensor, height: int, width: int, num_frames: int, patch_size: int = 2
) -> Tensor:
    D = int(os.environ.get("AE_SPATIAL_COMPRESSION", 16))
    return rearrange(
        x,
        "b (t h w) (c ph pw) -> b c t (h ph) (w pw)",
        h=math.ceil(height / D),
        w=math.ceil(width / D),
        t=num_frames,
        ph=patch_size,
        pw=patch_size,
    )


# ======================================================
# Prepare
# ======================================================


def prepare(
    t5,
    clip: HFEmbedder,
    img: Tensor,
    prompt: str | list[str],
    seq_align: int = 1,
    patch_size: int = 2,
) -> dict[str, Tensor]:
    """
    Prepare the input for the model.
    Args:
        t5 (HFEmbedder): The T5 model.
        clip (HFEmbedder): The CLIP model.
        img (Tensor): The image tensor.
        prompt (str | list[str]): The prompt(s).
    Returns:
        dict[str, Tensor]: The input dictionary.
        img_ids: used for positional embedding in T,H,W dimensions later
        text_ids: for positional embedding, but set to 0 for now since our text encoder already encodes positional information
    """
    bs, c, t, h, w = img.shape
    device, dtype = img.device, img.dtype
    if isinstance(prompt, str):
        prompt = [prompt]
    if bs != len(prompt):
        bs = len(prompt)

    img = rearrange(
        img, "b c t (h ph) (w pw) -> b (t h w) (c ph pw)", ph=patch_size, pw=patch_size
    )
    if img.shape[0] != bs:
        img = repeat(img, "b ... -> (repeat b) ...", repeat=bs // img.shape[0])

    img_ids = torch.zeros(t, h // patch_size, w // patch_size, 3)
    img_ids[..., 0] = img_ids[..., 0] + torch.arange(t)[:, None, None]
    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // patch_size)[None, :, None]
    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // patch_size)[None, None, :]
    img_ids = repeat(img_ids, "t h w c -> b (t h w) c", b=bs)

    # Encode the tokenized prompts
    txt = t5(prompt, added_tokens=img_ids.shape[1], seq_align=seq_align)
    if txt.shape[0] == 1 and bs > 1:
        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
    txt_ids = torch.zeros(bs, txt.shape[1], 3)

    vec = clip(prompt)
    if vec.shape[0] == 1 and bs > 1:
        vec = repeat(vec, "1 ... -> bs ...", bs=bs)

    return {
        "img": img,
        "img_ids": img_ids.to(device, dtype),
        "txt": txt.to(device, dtype),
        "txt_ids": txt_ids.to(device, dtype),
        "y_vec": vec.to(device, dtype),
    }


def prepare_ids(
    img: Tensor,
    t5_embedding: Tensor,
    clip_embedding: Tensor,
) -> dict[str, Tensor]:
    """
    Prepare the input for the model.
    Args:
        img (Tensor): The image tensor.
        t5_embedding (Tensor): The T5 embedding.
        clip_embedding (Tensor): The CLIP embedding.
    Returns:
        dict[str, Tensor]: The input dictionary.
        img_ids: used for positional embedding in T,H,W dimensions later
        text_ids: for positional embedding, but set to 0 for now since our text encoder already encodes positional information
    """
    bs, c, t, h, w = img.shape
    device, dtype = img.device, img.dtype

    img = rearrange(img, "b c t (h ph) (w pw) -> b (t h w) (c ph pw)", ph=2, pw=2)
    if img.shape[0] != bs:
        img = repeat(img, "b ... -> (repeat b) ...", repeat=bs // img.shape[0])

    img_ids = torch.zeros(t, h // 2, w // 2, 3)
    img_ids[..., 0] = img_ids[..., 0] + torch.arange(t)[:, None, None]
    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[None, :, None]
    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, None, :]
    img_ids = repeat(img_ids, "t h w c -> b (t h w) c", b=bs)

    # Encode the tokenized prompts
    if t5_embedding.shape[0] == 1 and bs > 1:
        t5_embedding = repeat(t5_embedding, "1 ... -> bs ...", bs=bs)
    txt_ids = torch.zeros(bs, t5_embedding.shape[1], 3)

    if clip_embedding.shape[0] == 1 and bs > 1:
        clip_embedding = repeat(clip_embedding, "1 ... -> bs ...", bs=bs)

    return {
        "img": img,
        "img_ids": img_ids.to(device, dtype),
        "txt": t5_embedding.to(device, dtype),
        "txt_ids": txt_ids.to(device, dtype),
        "y_vec": clip_embedding.to(device, dtype),
    }


def prepare_models(
    cfg: Config,
    device: torch.device,
    dtype: torch.dtype,
    offload_model: bool = False,
) -> tuple[nn.Module, nn.Module, nn.Module, nn.Module, dict[str, nn.Module]]:
    """
    Prepare models for inference.
    Args:
        cfg (Config): The configuration object.
        device (torch.device): The device to use.
        dtype (torch.dtype): The data type to use.
    Returns:
        tuple[nn.Module, nn.Module, nn.Module, nn.Module, dict[str, nn.Module]]: The models. They are the diffusion model, the autoencoder model, the T5 model, the CLIP model, and the optional models.
    """
    model_device = (
        "cpu" if offload_model and cfg.get("img_flux", None) is not None else device
    )

    model = build_module(
        cfg.model, MODELS, device_map=model_device, torch_dtype=dtype
    ).eval()
    model_ae = build_module(
        cfg.ae, MODELS, device_map=model_device, torch_dtype=dtype
    ).eval()
    model_t5 = build_module(cfg.t5, MODELS, device_map=device, torch_dtype=dtype).eval()
    model_clip = build_module(
        cfg.clip, MODELS, device_map=device, torch_dtype=dtype
    ).eval()
    if cfg.get("pretrained_lora_path", None) is not None:
        model = PeftModel.from_pretrained(
            model, cfg.pretrained_lora_path, is_trainable=False
        )

    # optional models
    optional_models = {}
    if cfg.get("img_flux", None) is not None:
        model_img_flux = build_module(
            cfg.img_flux, MODELS, device_map=device, torch_dtype=dtype
        ).eval()
        model_ae_img_flux = build_module(
            cfg.img_flux_ae, MODELS, device_map=device, torch_dtype=dtype
        ).eval()
        optional_models["img_flux"] = model_img_flux
        optional_models["img_flux_ae"] = model_ae_img_flux

    return model, model_ae, model_t5, model_clip, optional_models


def prepare_api(
    model: nn.Module,
    model_ae: nn.Module,
    model_t5: nn.Module,
    model_clip: nn.Module,
    optional_models: dict[str, nn.Module],
) -> callable:
    """
    Prepare the API function for inference.
    Args:
        model (nn.Module): The diffusion model.
        model_ae (nn.Module): The autoencoder model.
        model_t5 (nn.Module): The T5 model.
        model_clip (nn.Module): The CLIP model.
    Returns:
        callable: The API function for inference.
    """

    @torch.inference_mode()
    def api_fn(
        opt: SamplingOption,
        cond_type: str = "t2v",
        seed: int = None,
        sigma_min: float = 1e-5,
        text: list[str] = None,
        neg: list[str] = None,
        patch_size: int = 2,
        channel: int = 16,
        **kwargs,
    ):
        """
        The API function for inference.
        Args:
            opt (SamplingOption): The sampling options.
            text (list[str], optional): The text prompts. Defaults to None.
            neg (list[str], optional): The negative text prompts. Defaults to None.
        Returns:
            torch.Tensor: The generated images.
        """
        device = next(model.parameters()).device
        dtype = next(model.parameters()).dtype

        # passing seed will overwrite opt seed
        if seed is None:
            # random seed if not provided
            seed = opt.seed if opt.seed is not None else random.randint(0, 2**32 - 1)
        if opt.is_causal_vae:
            num_frames = (
                1
                if opt.num_frames == 1
                else (opt.num_frames - 1) // opt.temporal_reduction + 1
            )
        else:
            num_frames = (
                1 if opt.num_frames == 1 else opt.num_frames // opt.temporal_reduction
            )

        z = get_noise(
            len(text),
            opt.height,
            opt.width,
            num_frames,
            device,
            dtype,
            seed,
            patch_size=patch_size,
            channel=channel // (patch_size**2),
        )
        denoiser = SamplingMethodDict[opt.method]

        # i2v reference conditions
        references = [None] * len(text)
        if cond_type != "t2v" and "ref" in kwargs:
            reference_path_list = kwargs.pop("ref")
            references = collect_references_batch(
                reference_path_list,
                cond_type,
                model_ae,
                (opt.height, opt.width),
                is_causal=opt.is_causal_vae,
            )
        elif cond_type != "t2v":
            print(
                "your csv file doesn't have a ref column or is not processed properly. will default to cond_type t2v!"
            )
            cond_type = "t2v"

        # timestep editing
        timesteps = get_schedule(
            opt.num_steps,
            (z.shape[-1] * z.shape[-2]) // patch_size**2,
            num_frames,
            shift=opt.shift,
            shift_alpha=opt.flow_shift,
        )

        # prepare classifier-free guidance data (method specific)
        text, additional_inp = denoiser.prepare_guidance(
            text=text,
            optional_models=optional_models,
            device=device,
            dtype=dtype,
            neg=neg,
            guidance_img=opt.guidance_img,
        )

        inp = prepare(model_t5, model_clip, z, prompt=text, patch_size=patch_size)
        inp.update(additional_inp)

        if opt.method in [SamplingMethod.I2V]:
            # prepare references
            masks, masked_ref = prepare_inference_condition(
                z, cond_type, ref_list=references, causal=opt.is_causal_vae
            )
            inp["masks"] = masks
            inp["masked_ref"] = masked_ref
            inp["sigma_min"] = sigma_min

        x = denoiser.denoise(
            model,
            **inp,
            timesteps=timesteps,
            guidance=opt.guidance,
            text_osci=opt.text_osci,
            image_osci=opt.image_osci,
            scale_temporal_osci=(
                opt.scale_temporal_osci and "i2v" in cond_type
            ),  # don't use temporal osci for v2v or t2v
            flow_shift=opt.flow_shift,
            patch_size=patch_size,
        )

        x = unpack(x, opt.height, opt.width, num_frames, patch_size=patch_size)

        # replace for image condition
        if cond_type == "i2v_head":
            x[0, :, :1] = references[0][0]
        elif cond_type == "i2v_tail":
            x[0, :, -1:] = references[0][0]
        elif cond_type == "i2v_loop":
            x[0, :, :1] = references[0][0]
            x[0, :, -1:] = references[0][1]

        x = model_ae.decode(x)
        x = x[:, :, : opt.num_frames]  # image

        # remove the duplicate frames
        if not opt.is_causal_vae:
            if cond_type == "i2v_head":
                pad_len = model_ae.compression[0] - 1
                x = x[:, :, pad_len:]
            elif cond_type == "i2v_tail":
                pad_len = model_ae.compression[0] - 1
                x = x[:, :, :-pad_len]
            elif cond_type == "i2v_loop":
                pad_len = model_ae.compression[0] - 1
                x = x[:, :, pad_len:-pad_len]

        return x

    return api_fn

	import math
	import os
	import random
	from abc import ABC, abstractmethod
	from dataclasses import dataclass, replace

	import torch
	from einops import rearrange, repeat
	from mmengine.config import Config
	from peft import PeftModel
	from torch import Tensor, nn

	from opensora.datasets.aspect import get_image_size
	from opensora.models.mmdit.model import MMDiTModel
	from opensora.models.text.conditioner import HFEmbedder
	from opensora.registry import MODELS, build_module
	from opensora.utils.inference import (
	SamplingMethod,
	collect_references_batch,
	prepare_inference_condition,
	)

	# ======================================================
	# Sampling Options
	# ======================================================


	@dataclass
	class SamplingOption:
	# The width of the image/video.
	width: int \| None = None

	# The height of the image/video.
	height: int \| None = None

	# The resolution of the image/video. If provided, it will override the height and width.
	resolution: str \| None = None

	# The aspect ratio of the image/video. If provided, it will override the height and width.
	aspect_ratio: str \| None = None

	# The number of frames.
	num_frames: int = 1

	# The number of sampling steps.
	num_steps: int = 50

	# The classifier-free guidance (text).
	guidance: float = 4.0

	# use oscillation for text guidance
	text_osci: bool = False

	# The classifier-free guidance (image), or for the guidance on condition for i2v and v2v
	guidance_img: float \| None = None

	# use oscillation for image guidance
	image_osci: bool = False

	# use temporal scaling for image guidance
	scale_temporal_osci: bool = False

	# The seed for the random number generator.
	seed: int \| None = None

	# Whether to shift the schedule.
	shift: bool = True

	# The sampling method.
	method: str \| SamplingMethod = SamplingMethod.I2V

	# Temporal reduction
	temporal_reduction: int = 1

	# is causal vae
	is_causal_vae: bool = False

	# flow shift
	flow_shift: float \| None = None


	def sanitize_sampling_option(sampling_option: SamplingOption) -> SamplingOption:
	"""
	Sanitize the sampling options.

	Args:
	sampling_option (SamplingOption): The sampling options.

	Returns:
	SamplingOption: The sanitized sampling options.
	"""
	if (
	sampling_option.resolution is not None
	or sampling_option.aspect_ratio is not None
	):
	assert (
	sampling_option.resolution is not None
	and sampling_option.aspect_ratio is not None
	), "Both resolution and aspect ratio must be provided"
	resolution = sampling_option.resolution
	aspect_ratio = sampling_option.aspect_ratio
	height, width = get_image_size(resolution, aspect_ratio, training=False)
	else:
	assert (
	sampling_option.height is not None and sampling_option.width is not None
	), "Both height and width must be provided"
	height, width = sampling_option.height, sampling_option.width

	height = (height // 16 + (1 if height % 16 else 0)) * 16
	width = (width // 16 + (1 if width % 16 else 0)) * 16
	replace_dict = dict(height=height, width=width)

	if isinstance(sampling_option.method, str):
	method = SamplingMethod(sampling_option.method)
	replace_dict["method"] = method

	return replace(sampling_option, **replace_dict)


	def get_oscillation_gs(guidance_scale: float, i: int, force_num=10):
	"""
	get oscillation guidance for cfg.

	Args:
	guidance_scale: original guidance value
	i: denoising step
	force_num: before which don't apply oscillation
	"""
	if i < force_num or (i >= force_num and i % 2 == 0):
	gs = guidance_scale
	else:
	gs = 1.0
	return gs


	# ======================================================
	# Denoising
	# ======================================================


	class Denoiser(ABC):
	@abstractmethod
	def denoise(self, model: MMDiTModel, **kwargs) -> Tensor:
	"""Denoise the input."""

	@abstractmethod
	def prepare_guidance(
	self,
	text: list[str],
	optional_models: dict[str, nn.Module],
	device: torch.device,
	dtype: torch.dtype,
	**kwargs,
	) -> dict[str, Tensor]:
	"""Prepare the guidance for the model. This method will alter text."""


	class I2VDenoiser(Denoiser):
	def denoise(self, model: MMDiTModel, **kwargs) -> Tensor:
	img = kwargs.pop("img")
	timesteps = kwargs.pop("timesteps")
	guidance = kwargs.pop("guidance")
	guidance_img = kwargs.pop("guidance_img")

	# cond ref arguments
	masks = kwargs.pop("masks")
	masked_ref = kwargs.pop("masked_ref")
	kwargs.pop("sigma_min")

	# oscillation guidance
	text_osci = kwargs.pop("text_osci", False)
	image_osci = kwargs.pop("image_osci", False)
	scale_temporal_osci = kwargs.pop("scale_temporal_osci", False)

	# patch size
	patch_size = kwargs.pop("patch_size", 2)

	guidance_vec = torch.full(
	(img.shape[0],), guidance, device=img.device, dtype=img.dtype
	)
	for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
	# timesteps
	t_vec = torch.full(
	(img.shape[0],), t_curr, dtype=img.dtype, device=img.device
	)
	b, c, t, w, h = masked_ref.size()
	cond = torch.cat((masks, masked_ref), dim=1)
	cond = pack(cond, patch_size=patch_size)
	kwargs["cond"] = torch.cat([cond, cond, torch.zeros_like(cond)], dim=0)

	# forward preparation
	cond_x = img[: len(img) // 3]

	img = torch.cat([cond_x, cond_x, cond_x], dim=0)
	# forward
	pred = model(
	img=img,
	**kwargs,
	timesteps=t_vec,
	guidance=guidance_vec,
	)

	# prepare guidance
	text_gs = get_oscillation_gs(guidance, i) if text_osci else guidance
	image_gs = (
	get_oscillation_gs(guidance_img, i) if image_osci else guidance_img
	)
	cond, uncond, uncond_2 = pred.chunk(3, dim=0)
	if image_gs > 1.0 and scale_temporal_osci:
	# image_gs decrease with each denoising step
	step_upper_image_gs = torch.linspace(image_gs, 1.0, len(timesteps))[i]
	# image_gs increase along the temporal axis of the latent video
	image_gs = torch.linspace(1.0, step_upper_image_gs, t)[
	None, None, :, None, None
	].repeat(b, c, 1, h, w)
	image_gs = pack(image_gs, patch_size=patch_size).to(cond.device, cond.dtype)

	# update
	pred = uncond_2 + image_gs * (uncond - uncond_2) + text_gs * (cond - uncond)
	pred = torch.cat([pred, pred, pred], dim=0)

	img = img + (t_prev - t_curr) * pred

	img = img[: len(img) // 3]

	return img

	def prepare_guidance(
	self,
	text: list[str],
	optional_models: dict[str, nn.Module],
	device: torch.device,
	dtype: torch.dtype,
	**kwargs,
	) -> tuple[list[str], dict[str, Tensor]]:
	ret = {}

	neg = kwargs.get("neg", None)
	ret["guidance_img"] = kwargs.pop("guidance_img")

	# text
	if neg is None:
	neg = [""] * len(text)
	text = text + neg + neg
	return text, ret


	class DistilledDenoiser(Denoiser):
	def denoise(self, model: MMDiTModel, **kwargs) -> Tensor:
	img = kwargs.pop("img")
	timesteps = kwargs.pop("timesteps")
	guidance = kwargs.pop("guidance")

	guidance_vec = torch.full(
	(img.shape[0],), guidance, device=img.device, dtype=img.dtype
	)
	for t_curr, t_prev in zip(timesteps[:-1], timesteps[1:]):
	# timesteps
	t_vec = torch.full(
	(img.shape[0],), t_curr, dtype=img.dtype, device=img.device
	)
	# forward
	pred = model(
	img=img,
	**kwargs,
	timesteps=t_vec,
	guidance=guidance_vec,
	)
	# update
	img = img + (t_prev - t_curr) * pred
	return img

	def prepare_guidance(
	self,
	text: list[str],
	optional_models: dict[str, nn.Module],
	device: torch.device,
	dtype: torch.dtype,
	**kwargs,
	) -> tuple[list[str], dict[str, Tensor]]:
	return text, {}


	SamplingMethodDict = {
	SamplingMethod.I2V: I2VDenoiser(),
	SamplingMethod.DISTILLED: DistilledDenoiser(),
	}


	# ======================================================
	# Timesteps
	# ======================================================


	def time_shift(alpha: float, t: Tensor) -> Tensor:
	return alpha * t / (1 + (alpha - 1) * t)


	def get_res_lin_function(
	x1: float = 256, y1: float = 1, x2: float = 4096, y2: float = 3
	) -> callable:
	m = (y2 - y1) / (x2 - x1)
	b = y1 - m * x1
	return lambda x: m * x + b


	def get_schedule(
	num_steps: int,
	image_seq_len: int,
	num_frames: int,
	shift_alpha: float \| None = None,
	base_shift: float = 1,
	max_shift: float = 3,
	shift: bool = True,
	) -> list[float]:
	# extra step for zero
	timesteps = torch.linspace(1, 0, num_steps + 1)

	# shifting the schedule to favor high timesteps for higher signal images
	if shift:
	if shift_alpha is None:
	# estimate mu based on linear estimation between two points
	# spatial scale
	shift_alpha = get_res_lin_function(y1=base_shift, y2=max_shift)(
	image_seq_len
	)
	# temporal scale
	shift_alpha *= math.sqrt(num_frames)
	# calculate shifted timesteps
	timesteps = time_shift(shift_alpha, timesteps)

	return timesteps.tolist()


	def get_noise(
	num_samples: int,
	height: int,
	width: int,
	num_frames: int,
	device: torch.device,
	dtype: torch.dtype,
	seed: int,
	patch_size: int = 2,
	channel: int = 16,
	) -> Tensor:
	"""
	Generate a noise tensor.

	Args:
	num_samples (int): Number of samples.
	height (int): Height of the noise tensor.
	width (int): Width of the noise tensor.
	num_frames (int): Number of frames.
	device (torch.device): Device to put the noise tensor on.
	dtype (torch.dtype): Data type of the noise tensor.
	seed (int): Seed for the random number generator.

	Returns:
	Tensor: The noise tensor.
	"""
	D = int(os.environ.get("AE_SPATIAL_COMPRESSION", 16))
	return torch.randn(
	num_samples,
	channel,
	num_frames,
	# allow for packing
	patch_size * math.ceil(height / D),
	patch_size * math.ceil(width / D),
	device=device,
	dtype=dtype,
	generator=torch.Generator(device=device).manual_seed(seed),
	)


	def pack(x: Tensor, patch_size: int = 2) -> Tensor:
	return rearrange(
	x, "b c t (h ph) (w pw) -> b (t h w) (c ph pw)", ph=patch_size, pw=patch_size
	)


	def unpack(
	x: Tensor, height: int, width: int, num_frames: int, patch_size: int = 2
	) -> Tensor:
	D = int(os.environ.get("AE_SPATIAL_COMPRESSION", 16))
	return rearrange(
	x,
	"b (t h w) (c ph pw) -> b c t (h ph) (w pw)",
	h=math.ceil(height / D),
	w=math.ceil(width / D),
	t=num_frames,
	ph=patch_size,
	pw=patch_size,
	)


	# ======================================================
	# Prepare
	# ======================================================


	def prepare(
	t5,
	clip: HFEmbedder,
	img: Tensor,
	prompt: str \| list[str],
	seq_align: int = 1,
	patch_size: int = 2,
	) -> dict[str, Tensor]:
	"""
	Prepare the input for the model.

	Args:
	t5 (HFEmbedder): The T5 model.
	clip (HFEmbedder): The CLIP model.
	img (Tensor): The image tensor.
	prompt (str \| list[str]): The prompt(s).

	Returns:
	dict[str, Tensor]: The input dictionary.

	img_ids: used for positional embedding in T,H,W dimensions later
	text_ids: for positional embedding, but set to 0 for now since our text encoder already encodes positional information
	"""
	bs, c, t, h, w = img.shape
	device, dtype = img.device, img.dtype
	if isinstance(prompt, str):
	prompt = [prompt]
	if bs != len(prompt):
	bs = len(prompt)

	img = rearrange(
	img, "b c t (h ph) (w pw) -> b (t h w) (c ph pw)", ph=patch_size, pw=patch_size
	)
	if img.shape[0] != bs:
	img = repeat(img, "b ... -> (repeat b) ...", repeat=bs // img.shape[0])

	img_ids = torch.zeros(t, h // patch_size, w // patch_size, 3)
	img_ids[..., 0] = img_ids[..., 0] + torch.arange(t)[:, None, None]
	img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // patch_size)[None, :, None]
	img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // patch_size)[None, None, :]
	img_ids = repeat(img_ids, "t h w c -> b (t h w) c", b=bs)

	# Encode the tokenized prompts
	txt = t5(prompt, added_tokens=img_ids.shape[1], seq_align=seq_align)
	if txt.shape[0] == 1 and bs > 1:
	txt = repeat(txt, "1 ... -> bs ...", bs=bs)
	txt_ids = torch.zeros(bs, txt.shape[1], 3)

	vec = clip(prompt)
	if vec.shape[0] == 1 and bs > 1:
	vec = repeat(vec, "1 ... -> bs ...", bs=bs)

	return {
	"img": img,
	"img_ids": img_ids.to(device, dtype),
	"txt": txt.to(device, dtype),
	"txt_ids": txt_ids.to(device, dtype),
	"y_vec": vec.to(device, dtype),
	}


	def prepare_ids(
	img: Tensor,
	t5_embedding: Tensor,
	clip_embedding: Tensor,
	) -> dict[str, Tensor]:
	"""
	Prepare the input for the model.

	Args:
	img (Tensor): The image tensor.
	t5_embedding (Tensor): The T5 embedding.
	clip_embedding (Tensor): The CLIP embedding.

	Returns:
	dict[str, Tensor]: The input dictionary.

	img_ids: used for positional embedding in T,H,W dimensions later
	text_ids: for positional embedding, but set to 0 for now since our text encoder already encodes positional information
	"""
	bs, c, t, h, w = img.shape
	device, dtype = img.device, img.dtype

	img = rearrange(img, "b c t (h ph) (w pw) -> b (t h w) (c ph pw)", ph=2, pw=2)
	if img.shape[0] != bs:
	img = repeat(img, "b ... -> (repeat b) ...", repeat=bs // img.shape[0])

	img_ids = torch.zeros(t, h // 2, w // 2, 3)
	img_ids[..., 0] = img_ids[..., 0] + torch.arange(t)[:, None, None]
	img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[None, :, None]
	img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, None, :]
	img_ids = repeat(img_ids, "t h w c -> b (t h w) c", b=bs)

	# Encode the tokenized prompts
	if t5_embedding.shape[0] == 1 and bs > 1:
	t5_embedding = repeat(t5_embedding, "1 ... -> bs ...", bs=bs)
	txt_ids = torch.zeros(bs, t5_embedding.shape[1], 3)

	if clip_embedding.shape[0] == 1 and bs > 1:
	clip_embedding = repeat(clip_embedding, "1 ... -> bs ...", bs=bs)

	return {
	"img": img,
	"img_ids": img_ids.to(device, dtype),
	"txt": t5_embedding.to(device, dtype),
	"txt_ids": txt_ids.to(device, dtype),
	"y_vec": clip_embedding.to(device, dtype),
	}


	def prepare_models(
	cfg: Config,
	device: torch.device,
	dtype: torch.dtype,
	offload_model: bool = False,
	) -> tuple[nn.Module, nn.Module, nn.Module, nn.Module, dict[str, nn.Module]]:
	"""
	Prepare models for inference.

	Args:
	cfg (Config): The configuration object.
	device (torch.device): The device to use.
	dtype (torch.dtype): The data type to use.

	Returns:
	tuple[nn.Module, nn.Module, nn.Module, nn.Module, dict[str, nn.Module]]: The models. They are the diffusion model, the autoencoder model, the T5 model, the CLIP model, and the optional models.
	"""
	model_device = (
	"cpu" if offload_model and cfg.get("img_flux", None) is not None else device
	)

	model = build_module(
	cfg.model, MODELS, device_map=model_device, torch_dtype=dtype
	).eval()
	model_ae = build_module(
	cfg.ae, MODELS, device_map=model_device, torch_dtype=dtype
	).eval()
	model_t5 = build_module(cfg.t5, MODELS, device_map=device, torch_dtype=dtype).eval()
	model_clip = build_module(
	cfg.clip, MODELS, device_map=device, torch_dtype=dtype
	).eval()
	if cfg.get("pretrained_lora_path", None) is not None:
	model = PeftModel.from_pretrained(
	model, cfg.pretrained_lora_path, is_trainable=False
	)

	# optional models
	optional_models = {}
	if cfg.get("img_flux", None) is not None:
	model_img_flux = build_module(
	cfg.img_flux, MODELS, device_map=device, torch_dtype=dtype
	).eval()
	model_ae_img_flux = build_module(
	cfg.img_flux_ae, MODELS, device_map=device, torch_dtype=dtype
	).eval()
	optional_models["img_flux"] = model_img_flux
	optional_models["img_flux_ae"] = model_ae_img_flux

	return model, model_ae, model_t5, model_clip, optional_models


	def prepare_api(
	model: nn.Module,
	model_ae: nn.Module,
	model_t5: nn.Module,
	model_clip: nn.Module,
	optional_models: dict[str, nn.Module],
	) -> callable:
	"""
	Prepare the API function for inference.

	Args:
	model (nn.Module): The diffusion model.
	model_ae (nn.Module): The autoencoder model.
	model_t5 (nn.Module): The T5 model.
	model_clip (nn.Module): The CLIP model.

	Returns:
	callable: The API function for inference.
	"""

	@torch.inference_mode()
	def api_fn(
	opt: SamplingOption,
	cond_type: str = "t2v",
	seed: int = None,
	sigma_min: float = 1e-5,
	text: list[str] = None,
	neg: list[str] = None,
	patch_size: int = 2,
	channel: int = 16,
	**kwargs,
	):
	"""
	The API function for inference.

	Args:
	opt (SamplingOption): The sampling options.
	text (list[str], optional): The text prompts. Defaults to None.
	neg (list[str], optional): The negative text prompts. Defaults to None.

	Returns:
	torch.Tensor: The generated images.
	"""
	device = next(model.parameters()).device
	dtype = next(model.parameters()).dtype

	# passing seed will overwrite opt seed
	if seed is None:
	# random seed if not provided
	seed = opt.seed if opt.seed is not None else random.randint(0, 2**32 - 1)
	if opt.is_causal_vae:
	num_frames = (
	1
	if opt.num_frames == 1
	else (opt.num_frames - 1) // opt.temporal_reduction + 1
	)
	else:
	num_frames = (
	1 if opt.num_frames == 1 else opt.num_frames // opt.temporal_reduction
	)

	z = get_noise(
	len(text),
	opt.height,
	opt.width,
	num_frames,
	device,
	dtype,
	seed,
	patch_size=patch_size,
	channel=channel // (patch_size**2),
	)
	denoiser = SamplingMethodDict[opt.method]

	# i2v reference conditions
	references = [None] * len(text)
	if cond_type != "t2v" and "ref" in kwargs:
	reference_path_list = kwargs.pop("ref")
	references = collect_references_batch(
	reference_path_list,
	cond_type,
	model_ae,
	(opt.height, opt.width),
	is_causal=opt.is_causal_vae,
	)
	elif cond_type != "t2v":
	print(
	"your csv file doesn't have a ref column or is not processed properly. will default to cond_type t2v!"
	)
	cond_type = "t2v"

	# timestep editing
	timesteps = get_schedule(
	opt.num_steps,
	(z.shape[-1] * z.shape[-2]) // patch_size**2,
	num_frames,
	shift=opt.shift,
	shift_alpha=opt.flow_shift,
	)

	# prepare classifier-free guidance data (method specific)
	text, additional_inp = denoiser.prepare_guidance(
	text=text,
	optional_models=optional_models,
	device=device,
	dtype=dtype,
	neg=neg,
	guidance_img=opt.guidance_img,
	)

	inp = prepare(model_t5, model_clip, z, prompt=text, patch_size=patch_size)
	inp.update(additional_inp)

	if opt.method in [SamplingMethod.I2V]:
	# prepare references
	masks, masked_ref = prepare_inference_condition(
	z, cond_type, ref_list=references, causal=opt.is_causal_vae
	)
	inp["masks"] = masks
	inp["masked_ref"] = masked_ref
	inp["sigma_min"] = sigma_min

	x = denoiser.denoise(
	model,
	**inp,
	timesteps=timesteps,
	guidance=opt.guidance,
	text_osci=opt.text_osci,
	image_osci=opt.image_osci,
	scale_temporal_osci=(
	opt.scale_temporal_osci and "i2v" in cond_type
	), # don't use temporal osci for v2v or t2v
	flow_shift=opt.flow_shift,
	patch_size=patch_size,
	)

	x = unpack(x, opt.height, opt.width, num_frames, patch_size=patch_size)

	# replace for image condition
	if cond_type == "i2v_head":
	x[0, :, :1] = references[0][0]
	elif cond_type == "i2v_tail":
	x[0, :, -1:] = references[0][0]
	elif cond_type == "i2v_loop":
	x[0, :, :1] = references[0][0]
	x[0, :, -1:] = references[0][1]

	x = model_ae.decode(x)
	x = x[:, :, : opt.num_frames] # image

	# remove the duplicate frames
	if not opt.is_causal_vae:
	if cond_type == "i2v_head":
	pad_len = model_ae.compression[0] - 1
	x = x[:, :, pad_len:]
	elif cond_type == "i2v_tail":
	pad_len = model_ae.compression[0] - 1
	x = x[:, :, :-pad_len]
	elif cond_type == "i2v_loop":
	pad_len = model_ae.compression[0] - 1
	x = x[:, :, pad_len:-pad_len]

	return x

	return api_fn