Skip to content

vllm_omni.diffusion.models.magi2.pipeline_magi2

Native vLLM-Omni pipeline for sand-ai/MAGI-2-preview.

The model, packing, scheduler, checkpoint mapping, and decoder integration are owned by vLLM-Omni. SandAI's Apache-2.0 implementation was used as the accuracy reference, but is neither imported nor required at runtime.

This first native integration intentionally supports the released Preview stage (272p and 540p). The separately checkpointed 1080p refiner is not silently routed through an external implementation.

DEFAULT_NEGATIVE_PROMPT module-attribute 

DEFAULT_NEGATIVE_PROMPT = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"

MAGI2_AUDIO_SAMPLE_RATE module-attribute 

MAGI2_AUDIO_SAMPLE_RATE = (
    MAGI2_GENERATION_CONFIG.audio_sample_rate
)

MAGI2_MODEL_ID module-attribute 

MAGI2_MODEL_ID = 'sand-ai/MAGI-2-preview'

MAGI2_MODEL_REVISION module-attribute 

MAGI2_MODEL_REVISION = (
    "2dea51b64db47ee5b4402d36fd90829a0c58913b"
)

logger module-attribute 

logger = logging.getLogger(__name__)

Magi2Pipeline

Bases: Module, ProgressBarMixin, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, SupportImageInput, SupportAudioOutput

Native MAGI-2 Preview text/image-to-video-and-audio pipeline.

One pipeline instance is supported per worker process. Initialization sets process-wide deterministic state, so tests or deployments that need more than one pipeline must isolate them in separate worker processes.

audio_decoder instance-attribute 

audio_decoder = _Magi2StagedComponent(
    audio_decoder, self.device_str, pin_memory=pin_memory
)

audio_sample_rate class-attribute 

audio_sample_rate: int = MAGI2_AUDIO_SAMPLE_RATE

checkpoint_root instance-attribute 

checkpoint_root = _resolve_checkpoint_root(
    str(od_config.model), od_config.revision
)

data_proxy instance-attribute 

data_proxy = Magi2DataProxy()

deterministic instance-attribute 

deterministic = _env_flag(
    _config_value(
        od_config,
        "magi2_deterministic",
        "MAGI2_DETERMINISTIC",
    ),
    default=bool(
        getattr(od_config, "fa_deterministic", False)
    ),
)

device_str instance-attribute 

device_str = (
    f"cuda:{torch.accelerator.current_device_index()}"
)

dtype instance-attribute 

dtype = od_config.dtype or torch.bfloat16

dummy_run_num_frames class-attribute 

dummy_run_num_frames: int = 0

image_vae instance-attribute 

image_vae = _Magi2StagedComponent(
    image_vae, self.device_str, pin_memory=pin_memory
)

od_config instance-attribute 

od_config = od_config

sampler instance-attribute 

sampler = Magi2PreviewSampler(
    self.transformer,
    self.data_proxy,
    device=self.device_str,
    dtype=self.dtype,
)

support_audio_output class-attribute 

support_audio_output: bool = True

support_image_input class-attribute 

support_image_input: bool = True

text_encoder instance-attribute 

text_encoder = _Magi2StagedComponent(
    text_encoder, self.device_str, pin_memory=pin_memory
)

transformer instance-attribute 

transformer = Magi2PreviewTransformer(MAGI2_PREVIEW_CONFIG)

vae property 

vae: Module

Expose TurboVAE through the shared distributed-VAE contract.

video_decoder instance-attribute 

video_decoder = _Magi2StagedComponent(
    video_decoder, self.device_str, pin_memory=pin_memory
)

video_processor instance-attribute 

video_processor = VideoProcessor(
    vae_scale_factor=2
    * MAGI2_GENERATION_CONFIG.video_vae_stride[1]
)

weights_sources instance-attribute 

weights_sources = [
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=self.checkpoint_root,
        subfolder="preview",
        revision=None,
        prefix="transformer.",
        fall_back_to_pt=False,
    )
]

encode_prompt 

encode_prompt(prompt: str) -> Tensor

forward 

forward(req: DiffusionRequestBatch) -> DiffusionOutput

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

setup_compile 

setup_compile() -> None

Compile the transformer regions; the attention and MoE kernels stay eager.

get_magi2_post_process_func 

get_magi2_post_process_func(
    _od_config: OmniDiffusionConfig,
)