Skip to content

vllm_omni.diffusion.models.magi2.sampler_magi2

Native MAGI-2 Preview classifier-free-guidance sampling.

The CFG and denoising math is adapted from SandAI's Apache-2.0 MAGI-2 Preview inference implementation. Model placement is deliberately not managed here: vLLM-Omni's loader/offloader (including DLO) owns placement, while this sampler only prepares tensors and invokes the already placed model.

CFGConfig dataclass

Preview CFG controls, defaulted to the released MAGI-2 recipe.

audio_txt_guidance_scale class-attribute instance-attribute 

audio_txt_guidance_scale: float = 7.0

cfg_rescale class-attribute instance-attribute 

cfg_rescale: float = 0.0

cfg_trick_start_frame class-attribute instance-attribute 

cfg_trick_start_frame: int = 13

cfg_trick_value class-attribute instance-attribute 

cfg_trick_value: float = 2.0

dynamic_cfg_cutoff_value class-attribute instance-attribute 

dynamic_cfg_cutoff_value: float = 2.0

dynamic_cfg_start_t class-attribute instance-attribute 

dynamic_cfg_start_t: int = 500

skimmed_cfg_scale class-attribute instance-attribute 

skimmed_cfg_scale: float = 3.0

use_cfg_trick class-attribute instance-attribute 

use_cfg_trick: bool = False

use_dynamic_cfg class-attribute instance-attribute 

use_dynamic_cfg: bool = False

use_ref_for_uncond class-attribute instance-attribute 

use_ref_for_uncond: bool = True

use_skimmed_cfg_linear class-attribute instance-attribute 

use_skimmed_cfg_linear: bool = False

video_txt_guidance_scale class-attribute instance-attribute 

video_txt_guidance_scale: float = 5.0

Magi2PreviewSampler

Bases: CFGParallelMixin

Run MAGI-2 Preview's joint video/audio denoising loop.

data_proxy instance-attribute 

data_proxy = data_proxy or Magi2DataProxy()

device instance-attribute 

device = (
    torch.device(device) if device is not None else None
)

dtype instance-attribute 

dtype = dtype

model instance-attribute 

model = model

cfg_velocity 

cfg_velocity(
    model_output: tuple[Tensor, Tensor],
    video_txt_guidance_scale: Tensor | float,
    audio_txt_guidance_scale: Tensor | float,
    cfg_config: CFGConfig | None = None,
    latent: Tensor | None = None,
    audio_latent: Tensor | None = None,
) -> tuple[Tensor | None, Tensor | None]

combine_cfg_noise 

combine_cfg_noise(
    positive_noise_pred: Tensor | tuple[Tensor, ...],
    negative_noise_pred: Tensor | tuple[Tensor, ...],
    true_cfg_scale: float,
    cfg_normalize: bool = False,
    kwargs: dict[str, Any] | None = None,
) -> tuple[Tensor, Tensor]

Preserve MAGI-2's dual-modality guidance under shared CFG dispatch.

forward 

forward(model_input: ModelInput) -> tuple[Tensor, Tensor]

precalculate_cfg 

precalculate_cfg(
    t_list: Sequence[Tensor],
    latent_length: int,
    cfg_config: CFGConfig,
    *,
    device: device | str | None = None,
) -> tuple[list[Tensor], list[float]]

predict_noise 

predict_noise(
    model_input: ModelInput,
) -> tuple[Tensor, Tensor]

Run one conditional or unconditional branch for shared CFG dispatch.

prepare_model_input 

prepare_model_input(
    latent: Tensor,
    audio_latent: Tensor,
    txt_feat: Tensor,
    null_txt_feat: Tensor,
    ref_audio_feat: Tensor | None = None,
    ref_video_feat: Tensor | None = None,
    ref_image_feat: Tensor | None = None,
    ref_image_feat_len: Tensor | None = None,
    ref_image_special_token_embedding: Tensor | None = None,
    t: Tensor | float | None = None,
    cfg_config: CFGConfig | None = None,
) -> ModelInput

sample 

sample(
    sampler_input: SamplerInput,
) -> tuple[Tensor, Tensor]

step 

step(
    model_output: tuple[Tensor, Tensor],
    latent: Tensor,
    audio_latent: Tensor,
    video_txt_guidance_scale: Tensor | float,
    audio_txt_guidance_scale: Tensor | float,
    video_scheduler: SchedulerMixin,
    audio_scheduler: SchedulerMixin,
    t: Tensor,
    cfg_config: CFGConfig | None = None,
) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]

SamplerInput dataclass

audio_latent instance-attribute 

audio_latent: Tensor

audio_scheduler instance-attribute 

audio_scheduler: SchedulerMixin

audio_t_list instance-attribute 

audio_t_list: Sequence[Tensor]

cfg_config instance-attribute 

cfg_config: CFGConfig

latent instance-attribute 

latent: Tensor

null_txt_feat instance-attribute 

null_txt_feat: Tensor

ref_audio_feat instance-attribute 

ref_audio_feat: Tensor | None

ref_image_feat class-attribute instance-attribute 

ref_image_feat: Tensor | None = None

ref_image_feat_len class-attribute instance-attribute 

ref_image_feat_len: Tensor | None = None

ref_image_special_token_embedding class-attribute instance-attribute 

ref_image_special_token_embedding: Tensor | None = None

ref_video_feat instance-attribute 

ref_video_feat: Tensor | None

txt_feat instance-attribute 

txt_feat: Tensor

video_scheduler instance-attribute 

video_scheduler: SchedulerMixin

video_t_list instance-attribute 

video_t_list: Sequence[Tensor]

build_magi2_preview_schedulers 

build_magi2_preview_schedulers(
    num_inference_steps: int,
    *,
    device: device | str,
    shift: float = 7.0,
) -> tuple[
    FlowUniPCMultistepScheduler, FlowUniPCMultistepScheduler
]

Create independent, identically configured video/audio schedulers.