mstar.model.wan22.components.unipc

mstar.model.wan22.components.unipc#

Inline UniPC (bh2, order 2, flow prediction) for the Wan2.2 denoise loop.

Exact port of diffusers 0.39.0 UniPCMultistepScheduler, restricted to this checkpoint’s scheduler config. Any deviation from the reference is a bug, not a choice; the port is asserted in lockstep against it per step.

The reference scheduler keeps state across step() calls, but the denoise loop is stateless per iteration, so that state is split in two: per-request tables (make_unipc_tables), derivable from the step count and recomputed each iteration, and loop-carried tensors (UniPCState), routed through the graph’s loop-back edges. All solver math runs float32 on the sample’s device, as the reference does.

Functions

make_unipc_tables(num_inference_steps, ...)

Per-request sigma and timestep tables.

unipc_convert_model_output(model_output, ...)

Flow-prediction x0 conversion.

unipc_corrector_step(state, ...)

UniC correction of x_k using step k's converted output (source multistep_uni_c_bh_update).

unipc_effective_order(step_index, ...)

Predictor order at step k.

unipc_predictor_step(state, sample, sigmas, ...)

UniP update x_k -> x_{k+1}.

Classes

UniPCState(model_outputs, last_sample)

Loop-carried UniPC solver state for one request.
class mstar.model.wan22.components.unipc.UniPCState(model_outputs, last_sample)[source]#

Bases: object

Loop-carried UniPC solver state for one request.

model_outputs is the order-2 ring buffer of converted outputs (slot 1 is the previous step, slot 0 the one before). last_sample is the sample the previous predictor was given, consumed by this step’s corrector. Both are zero before they are written, and the order ramp guarantees they are not read then.

Parameters:
last_sample: Tensor#
model_outputs: Tensor#
mstar.model.wan22.components.unipc.make_unipc_tables(num_inference_steps, flow_shift)[source]#

Per-request sigma and timestep tables.

Returns sigmas float32 [N+1] with the terminal zero appended, and timesteps int64 [N]. Sigmas are computed float64 and cast at the end, and the timesteps are truncated by the int64 cast — both as the reference does, and both load-bearing for bit-exactness.

Parameters:

  • num_inference_steps (int)

  • flow_shift (float)

Return type:

tuple[Tensor, Tensor]

mstar.model.wan22.components.unipc.unipc_convert_model_output(model_output, sample, sigmas, step_index)[source]#

Flow-prediction x0 conversion.

sigma_t stays a CPU 0-dim tensor on purpose, as the reference leaves it. A CPU scalar takes the CUDA kernel’s full-precision Scalar path; a device 0-dim tensor would instead be type-promoted to the bf16 operand’s dtype before the multiply. That is a bitwise difference, and it breaks lockstep with the reference. (It is also why this node cannot be compiled — see Wan22DitSubmodule.)

Parameters:
Return type:

Tensor

mstar.model.wan22.components.unipc.unipc_corrector_step(state, this_model_output, this_sample, sigmas, step_index, order)[source]#

UniC correction of x_k using step k’s converted output (source multistep_uni_c_bh_update). Runs before the ring shift, so state.model_outputs[1] is step k-1’s output and [0] step k-2’s; state.last_sample is the sample step k-1’s predictor started from.

Parameters:
Return type:

Tensor

mstar.model.wan22.components.unipc.unipc_effective_order(step_index, num_inference_steps)[source]#

Predictor order at step k.

Capped both by the steps remaining (so the last step is order 1) and by the warmup ramp. Step k’s corrector reuses step k-1’s value.

Parameters:

  • step_index (int)

  • num_inference_steps (int)

Return type:

int

mstar.model.wan22.components.unipc.unipc_predictor_step(state, sample, sigmas, step_index, order)[source]#

UniP update x_k -> x_{k+1}.

state.model_outputs[1] must already hold step k’s converted output: the ring buffer is shifted before the predictor runs.

Parameters:
Return type:

Tensor