Function predict_flow

scripts/demo/app.py:82–159

(
    inputs: SampleState,
    model: ModelState,
    prediction_mode: Literal["parallel", "autoregressive"],
    flow_resolution: int = 256,
    ar_downsampling_factor: int = 8,
    enable_profiling: bool = False,
)

@torch.no_grad()
def predict_flow(
    inputs: SampleState,
    model: ModelState,
    prediction_mode: Literal["parallel", "autoregressive"],
    flow_resolution: int = 256,
    ar_downsampling_factor: int = 8,
    enable_profiling: bool = False,
) -> UInt8[torch.Tensor, "h w c"]:
    assert inputs.query is None, "Query point should not be specified for flow prediction."

    # Prepare inputs
    if len(inputs.pokes) > 0:
        pokes: Float[torch.Tensor, "b l t c"] = torch.tensor(inputs.pokes, device=model.device)[None]
    else:
        pokes = torch.empty((1, 0, 2, 2), dtype=torch.float32, device=model.device)
    poke_pos: Float[torch.Tensor, "b l c"] = pokes[:, :, 0, :]
    poke_flow: Float[torch.Tensor, "b l c"] = pokes[:, :, 1, :] - pokes[:, :, 0, :]
    query_pos: Float[torch.Tensor, "b l c"] = make_axial_pos_2d(flow_resolution, flow_resolution, device=model.device)[
        None
    ]

    with torch.autocast(device_type=model.device.type, dtype=torch.bfloat16):
        if prediction_mode == "parallel":
            start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
            start.record()
            pred: MixtureSameFamily = model.model.predict_parallel(
                poke_pos=poke_pos, poke_flow=poke_flow, query_pos=query_pos, camera_static=True, d_img=inputs.d_img
            )
            end.record()
            torch.cuda.synchronize()
            gr.Info(
                f"Delay: {start.elapsed_time(end):,.1f} ms ({math.prod(query_pos.shape[:-1]) / (start.elapsed_time(end) / 1e3):,.0f} queries / second)",
                duration=2,
            )
            flow: Float[torch.Tensor, "b h w c"] = rearrange(
                pred.mean, "b (h w) c -> b c h w", h=flow_resolution, w=flow_resolution
            )
        elif prediction_mode == "autoregressive":
            if ar_downsampling_factor != 1:
                ar_flow_resolution = flow_resolution // ar_downsampling_factor
                query_pos_ar = make_axial_pos_2d(
                    ar_flow_resolution,
                    ar_flow_resolution,
                    device=model.device,
                )[None]
            else:
                query_pos_ar = query_pos
            flow: Float[torch.Tensor, "b h w c"] = rearrange(
                model.model.predict_autoregressive(
                    poke_pos=poke_pos,
                    poke_flow=poke_flow,
                    query_pos=query_pos_ar,
                    camera_static=True,
                    d_img=inputs.d_img,
                    randomize_order=True,
                ),
                "b (h w) c -> b c h w",
                h=ar_flow_resolution,