Skip to content

vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2

DEBUG_PERF module-attribute 

DEBUG_PERF = False

WAN_SAMPLE_SOLVER_CHOICES module-attribute 

WAN_SAMPLE_SOLVER_CHOICES = {'unipc', 'euler'}

logger module-attribute 

logger = getLogger(__name__)

Wan22Pipeline

Bases: Module, PipelineParallelMixin, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin

boundary_ratio instance-attribute 

boundary_ratio = boundary_ratio

current_timestep property 

current_timestep

device instance-attribute 

device = get_local_device()

do_classifier_free_guidance property 

do_classifier_free_guidance

expand_timesteps instance-attribute 

expand_timesteps = get('expand_timesteps', False)

guidance_scale property 

guidance_scale

has_transformer_2 instance-attribute 

has_transformer_2 = transformer_2_info[0] is not None

num_timesteps property 

num_timesteps

od_config instance-attribute 

od_config = od_config

scheduler instance-attribute 

scheduler = build_wan_scheduler(_sample_solver, _flow_shift)

text_encoder instance-attribute 

text_encoder = to(device)

tokenizer instance-attribute 

tokenizer = from_pretrained_with_prefetch(
    from_pretrained,
    model,
    subfolder="tokenizer",
    prefetch_list=component_subfolders,
    local_files_only=local_files_only,
)

transformer instance-attribute 

transformer = _create_transformer(transformer_config)

transformer_2 instance-attribute 

transformer_2 = _create_transformer(transformer_2_config)

transformer_config instance-attribute 

transformer_config = config

vae instance-attribute 

vae = to(device)

vae_scale_factor_spatial instance-attribute 

vae_scale_factor_spatial = (
    scale_factor_spatial
    if getattr(self, "vae", None)
    else 8
)

vae_scale_factor_temporal instance-attribute 

vae_scale_factor_temporal = (
    scale_factor_temporal
    if getattr(self, "vae", None)
    else 4
)

weights_sources instance-attribute 

weights_sources = []

check_inputs 

check_inputs(
    prompt,
    negative_prompt,
    height,
    width,
    prompt_embeds=None,
    negative_prompt_embeds=None,
    guidance_scale_2=None,
    boundary_ratio=None,
)

diffuse 

diffuse(
    latents: Tensor,
    timesteps: Tensor,
    prompt_embeds: Tensor,
    negative_prompt_embeds: Tensor | None,
    guidance_low: float,
    guidance_high: float,
    boundary_timestep: float | None,
    dtype: dtype,
    attention_kwargs: dict[str, Any],
    latent_condition: Tensor | None = None,
    first_frame_mask: Tensor | None = None,
) -> Tensor | AsyncLatents

encode_prompt 

encode_prompt(
    prompt: str | list[str],
    negative_prompt: str | list[str] | None = None,
    do_classifier_free_guidance: bool = True,
    num_videos_per_prompt: int = 1,
    max_sequence_length: int = 512,
    device: device | None = None,
    dtype: dtype | None = None,
)

forward 

forward(
    req: OmniDiffusionRequest,
    prompt: str | None = None,
    negative_prompt: str | None = None,
    height: int = 480,
    width: int = 832,
    num_inference_steps: int = 40,
    guidance_scale: float | tuple[float, float] = 4.0,
    frame_num: int = 81,
    output_type: str | None = "np",
    generator: Generator | list[Generator] | None = None,
    prompt_embeds: Tensor | None = None,
    negative_prompt_embeds: Tensor | None = None,
    attention_kwargs: dict | None = None,
    **kwargs,
) -> DiffusionOutput

load_weights 

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

Load weights using AutoWeightsLoader for vLLM integration.

predict_noise 

predict_noise(
    current_model: Module | None = None, **kwargs: Any
) -> Tensor | IntermediateTensors

Forward pass through transformer to predict noise.

Parameters:

Name Type Description Default
current_model Module | None

The transformer model to use (transformer or transformer_2)

 None
**kwargs Any

Arguments to pass to the transformer

 {}

Returns:

Type Description
Tensor | IntermediateTensors

Predicted noise tensor or IntermediateTensors on non-last PP stages.

prepare_latents 

prepare_latents(
    batch_size: int,
    num_channels_latents: int,
    height: int,
    width: int,
    num_frames: int,
    dtype: dtype | None,
    device: device | None,
    generator: Generator | list[Generator] | None,
    latents: Tensor | None = None,
) -> Tensor

WanT2VDMD2Pipeline

Bases: DMD2PipelineMixin, Wan22Pipeline

Wan 2.x T2V pipeline for FastGen DMD2-distilled models.

build_wan_scheduler 

build_wan_scheduler(
    sample_solver: str, flow_shift: float
) -> Any

create_transformer_from_config 

create_transformer_from_config(
    config: dict,
    quant_config: QuantizationConfig | None = None,
    prefix: str = "",
) -> WanTransformer3DModel

Create WanTransformer3DModel from config dict.

get_wan22_post_process_func 

get_wan22_post_process_func(od_config: OmniDiffusionConfig)

get_wan22_pre_process_func 

get_wan22_pre_process_func(od_config: OmniDiffusionConfig)

Pre-process function for Wan2.2: optionally load and resize input image for I2V mode.

load_transformer_config 

load_transformer_config(
    model_path: str,
    subfolder: str = "transformer",
    local_files_only: bool = True,
) -> dict

Load transformer config from model directory or HF Hub.

resolve_wan_flow_shift 

resolve_wan_flow_shift(
    req: OmniDiffusionRequest,
    od_config: OmniDiffusionConfig,
) -> float

resolve_wan_sample_solver 

resolve_wan_sample_solver(
    req: OmniDiffusionRequest, default: str = "unipc"
) -> str

retrieve_latents 

retrieve_latents(
    encoder_output: Tensor,
    generator: Generator | None = None,
    sample_mode: str = "sample",
)

Retrieve latents from VAE encoder output.