Function pad_v_like_x

diff2flow/flow.py:25–32

Function to reshape the vector by the number of dimensions of x. E.g. x (bs, c, h, w), v (bs) -> v (bs, 1, 1, 1).

(v_, x_)

def pad_v_like_x(v_, x_):
    """
    Function to reshape the vector by the number of dimensions
    of x. E.g. x (bs, c, h, w), v (bs) -> v (bs, 1, 1, 1).
    """
    if isinstance(v_, float):
        return v_
    return v_.reshape(-1, *([1] * (x_.ndim - 1)))


def forward_with_cfg(x, t, model, cfg_scale=1.0, uc_cond=None, cond_key="y", **model_kwargs):