test_app_cpu.py

"""The Space's CPU path, end to end, with the GPU half stubbed.

Everything between the upload and the booking — decoding the clip onto 24 fps, turning a painted still into a mask,
choosing the crop and the canvas, and pasting the result back — runs on the host and is where a demo actually breaks.
The stub returns the crop it was handed, so a correct paste is the identity and any framing mistake shows up as a
difference against the source.
"""

import os
import sys
import types

import numpy as np

os.environ.setdefault("H3_MODEL_REPO", "hf-internal-testing/definitely-not-a-repo")

# `spaces` decorates with a real GPU booking and `ncii_guard` spawns a classifier subprocess; neither belongs in a
# host-side test. Stub both before `app` imports them.
spaces = types.ModuleType("spaces")
spaces.GPU = lambda *args, **kwargs: (lambda fn: fn)
sys.modules.setdefault("spaces", spaces)

guard = types.ModuleType("ncii_guard")
guard.start = lambda: None
guard.classify = lambda prompt, timeout=60.0: {"label": "safe", "score": 0.0}
sys.modules["ncii_guard"] = guard

import app  # noqa: E402


SOURCE = "examples/plate.mp4"


def test_load_error_is_the_only_failure():
    """The stub repo cannot load, and that has to surface as a message rather than an exception at import."""
    assert app.LOAD_ERROR is not None and "failed" in app.LOAD_ERROR
    print("startup: load failure captured, module still importable")


def test_read_clip_lands_on_24_fps():
    frames, waveform, sample_rate, source_fps, total = app.read_clip(SOURCE, 0.0, 5.0)
    assert frames.dtype == np.uint8 and frames.ndim == 4
    assert frames.shape[0] == round(5.0 * app.FPS), frames.shape
    assert waveform is not None and waveform.shape[0] == 2
    print(f"read_clip: {frames.shape} at {source_fps:g} fps -> 24 fps, audio {tuple(waveform.shape)} @ {sample_rate}")

    offset, _, _, _, _ = app.read_clip(SOURCE, 2.0, 2.0)
    assert offset.shape[0] == round(2.0 * app.FPS)
    assert not np.array_equal(offset[0], frames[0]), "a start offset has to move the first frame"
    print(f"read_clip: start offset honoured, {offset.shape[0]} frames")


def test_build_mask_from_painted_layers_and_from_composite():
    height, width = 64, 96
    background = np.zeros((height, width, 3), np.uint8)
    stroke = np.zeros((height, width, 4), np.uint8)
    stroke[10:30, 20:50, 3] = 255

    from_layers = app.build_mask({"layers": [stroke], "background": background}, None, 7)
    assert set(np.unique(from_layers)) <= {0.0, 1.0}, "masks are hard unless softening was asked for"
    assert from_layers.shape == (7, height, width)
    assert from_layers[0, 20, 30] == 1.0 and from_layers[0, 0, 0] == 0.0
    assert np.array_equal(from_layers[0], from_layers[-1]), "a painted mask is held for the whole clip"

    composite = background.copy()
    composite[10:30, 20:50] = 255
    from_composite = app.build_mask({"layers": [], "background": background, "composite": composite}, None, 7)
    assert np.array_equal(from_composite, from_layers), "the fallback has to find the same region"
    print("build_mask: layers and the composite fallback agree")

    for empty in ({"layers": [], "background": background, "composite": background.copy()}, None):
        try:
            app.build_mask(empty, None, 7)
        except Exception as error:
            assert "repaint" in str(error), error
        else:
            raise AssertionError("an empty mask must be refused")
    print("build_mask: an empty mask is refused")


def test_a_request_without_references_is_refused_before_it_costs_anything():
    """`ref2va` needs something to reference, and finding that out from the conditioner wastes a booking on both
    halves. Object removal from the prompt alone is the other partition's job, which this Space does not hold."""
    for references, expected in [
        ([], "at least one reference"),
        ([("audio", "examples/voice.wav")], "paired with an image"),
    ]:
        try:
            app.check_references(references)
        except Exception as error:
            assert expected in str(error), (expected, error)
        else:
            raise AssertionError(f"{references} should have been refused")
    app.check_references([("image", "examples/subject.png")])
    print("check_references: empty and audio-only requests refused, an image reference accepted")


def test_masks_are_hard_by_default():
    """A feathered edge leaves rows holding a mixture of source and repaint, which shows as a band along the mask."""
    ramp = np.tile(np.linspace(0, 1, 64, dtype=np.float32), (48, 1))[None]
    ramp = np.repeat(ramp, 5, axis=0)
    import tempfile, av, os

    path = os.path.join(tempfile.mkdtemp(), "m.mp4")
    c = av.open(path, mode="w"); st = c.add_stream("libx264", rate=24)
    st.width, st.height, st.pix_fmt = 64, 48, "yuv420p"; st.options = {"crf": "6"}
    for f in ramp:
        c.mux(st.encode(av.VideoFrame.from_ndarray(np.repeat((f*255).astype(np.uint8)[:, :, None], 3, 2), format="rgb24")))
    for pkt in st.encode(): c.mux(pkt)
    c.close()

    hard = app.build_mask(None, path, 5)
    assert set(np.unique(hard)) <= {0.0, 1.0}, np.unique(hard)[:5]
    soft = app.build_mask(None, path, 5, soften=True)
    assert len(np.unique(soft)) > 2, "softening should keep the ramp"
    print(f"hard mask: {len(np.unique(hard))} levels; soft: {len(np.unique(soft))} levels")


def test_canvas_label_is_one_the_conditioner_knows():
    for box, expected_ratio in [((0, 0, 544, 960), 960 / 544), ((0, 0, 768, 768), 1.0), ((0, 0, 1024, 768), 0.75)]:
        label = app.nearest_canvas_label(box[2], box[3])
        assert label in app.CANVASES, label
        height, width = app.CANVASES[label]
        assert abs(width / height - expected_ratio) < 0.06, (label, width / height, expected_ratio)
    # An off-grid crop still has to resolve to something the other half accepts.
    assert app.nearest_canvas_label(391, 613) in app.CANVASES
    print("nearest_canvas_label: every crop maps onto a label the conditioner knows")


def test_generate_frames_the_crop_and_pastes_it_back_unchanged():
    """The stub hands the crop straight back, so the paste is the identity everywhere the mask preserves."""
    calls = {}

    def stub(request, steps, seed):
        calls.update(request=request, steps=steps, seed=seed)
        import torch

        frames = request["source_video"]
        return [f for f in frames], torch.zeros(2, 32000), 32000

    original, load_error, pipe = app._generate, app.LOAD_ERROR, app.PIPE
    # The stub stands in for the pipeline, so the startup failure must not gate the request.
    app._generate, app.LOAD_ERROR, app.PIPE = stub, None, object()
    try:
        height, width = 768, 768
        background = np.zeros((height, width, 3), np.uint8)
        stroke = np.zeros((height, width, 4), np.uint8)
        stroke[300:520, 260:500, 3] = 255

        path, report = app.generate(
            SOURCE,
            "<Picture 1> the person from the picture",
            "examples/subject.png",
            None,
            painted={"layers": [stroke], "background": background},
            steps=12,
            seed=7,
        )
    finally:
        app._generate, app.LOAD_ERROR, app.PIPE = original, load_error, pipe

    request = calls["request"]
    print(f"report: {report}")
    assert calls["steps"] == 12 and calls["seed"] == 7
    assert request["height"] % 32 == 0 and request["width"] % 32 == 0
    assert request["source_video"].shape[1:3] == request["mask"].shape[1:3]
    assert request["num_frames"] % 17 == 5, request["num_frames"]
    assert request["source_audio"] is not None, "keep_audio defaults on"
    assert request["audio_mask"] is None, "a kept soundtrack is preserved whole"

    # The crop has to contain the whole painted region, and the canvas must not upscale a small box.
    box_h, box_w = request["source_video"].shape[1], request["source_video"].shape[2]
    assert request["mask"].sum() > 0
    assert box_h <= height and box_w <= width
    print(f"crop: {box_w}x{box_h} box -> {request['width']}x{request['height']} canvas, {request['num_frames']} frames")

    import av

    with av.open(path) as container:
        stream = container.streams.video[0]
        out = np.stack([f.to_ndarray(format="rgb24") for f in container.decode(video=0)])
        assert stream.height == height and stream.width == width, (stream.width, stream.height)
    source, _, _, _, _ = app.read_clip(SOURCE, 0.0, 5.0)
    # `generate` snaps up to the VAE's `17 * n + 5` grid by holding the last frame, so the output is a few
    # frames longer than the stretch that was asked for. Compare over the frames that exist in both.
    shared = min(out.shape[0], source.shape[0])
    out, source = out[:shared], source[:shared]
    # H.264 is lossy, so this is "the plate came back", not "bit identical".
    delta = np.abs(out.astype(np.int16) - source.astype(np.int16)).mean()
    assert delta < 4.0, f"the pasted-back plate drifted from the source by {delta:.2f}/255"
    print(f"paste_back: identity round trip through the muxer, mean |delta| = {delta:.2f}/255")


def test_a_silent_source_still_builds_a_request():
    """The dance example is the first plate with no audio track, and `keep_audio` is on by default.

    `read_clip` returns no waveform for it, so the request has to carry `source_audio=None` rather than crash or
    invent silence — the blocks treat that as "this clip has no soundtrack to preserve" and let the model write one.
    Every other example brings audio, which is exactly how the video-reference bug hid, so this path gets its own
    test rather than a shipped example nobody ran.
    """
    calls = {}

    def stub(request, steps, seed):
        calls.update(request=request)
        import torch

        return [f for f in request["source_video"]], torch.zeros(2, 32000), 32000

    silent = "examples/dance.mp4"
    frames, waveform, _, _, _ = app.read_clip(silent, 0.0, 5.0)
    assert waveform is None, "the dance clip is supposed to be silent"

    height, width = frames.shape[1], frames.shape[2]
    background = np.zeros((height, width, 3), np.uint8)
    stroke = np.zeros((height, width, 4), np.uint8)
    stroke[int(height * 0.15):int(height * 0.95), int(width * 0.2):int(width * 0.85), 3] = 255

    original, load_error, pipe = app._generate, app.LOAD_ERROR, app.PIPE
    app._generate, app.LOAD_ERROR, app.PIPE = stub, None, object()
    try:
        path, report = app.generate(
            silent, "a rusty humanoid robot dances in an empty warehouse", [], None, 7,
            painted={"layers": [stroke], "background": background}, steps=4,
        )
    finally:
        app._generate, app.LOAD_ERROR, app.PIPE = original, load_error, pipe

    request = calls["request"]
    assert request["source_audio"] is None, "a silent plate has no soundtrack to keep"
    assert request["audio_mask"] is None
    assert request["num_frames"] % 17 == 5, request["num_frames"]
    assert path and report
    print(f"silent source: {request['width']}x{request['height']} canvas, {request['num_frames']} frames, no audio")


def test_the_shipped_example_row_is_a_complete_request():
    """The example has to run with no brush at all: four columns, everything else on its defaults."""
    import inspect

    names = list(inspect.signature(app.generate).parameters)
    assert names[:5] == ["source_path", "prompt", "image_1", "mask_video", "seed"], names[:5]
    for name in names[5:]:
        assert inspect.signature(app.generate).parameters[name].default is not inspect.Parameter.empty, name

    for asset in ("examples/plate.mp4", "examples/plate_mask.mp4", "examples/subject.png",
                  "examples/platform.mp4", "examples/wednesday.png", "examples/motion.mp4"):
        assert os.path.exists(asset), asset
    # every example row must have as many columns as the inputs list it is applied to, or gradio fills the wrong
    # components without complaining
    import ast

    tree = ast.parse(open("app.py").read())
    call = next(
        n for n in ast.walk(tree)
        if isinstance(n, ast.Call) and getattr(n.func, "attr", None) == "Examples"
    )
    kwargs = {k.arg: k.value for k in call.keywords}
    columns = len(kwargs["inputs"].elts)
    rows = kwargs["examples"].elts
    assert rows, "no example rows found"
    for index, row in enumerate(rows):
        assert len(row.elts) == columns, f"row {index} has {len(row.elts)} fields for {columns} inputs"
    print(f"examples: {len(rows)} rows, {columns} columns each, all assets present")
    frames, _, _, _, _ = app.read_clip("examples/plate.mp4", 0.0, 5.0)
    mask = app.build_mask(None, "examples/plate_mask.mp4", frames.shape[0])
    assert mask.shape == frames.shape[:3], (mask.shape, frames.shape)
    assert 0.05 < mask.mean() < 0.85, mask.mean()
    # The mask has to move, or the example is really a static box and the mask-clip path is untested.
    assert not np.allclose(mask[0], mask[-1]), "the example mask should track the subject"
    print(f"example: plate {frames.shape}, mask covers {mask.mean()*100:.0f}% and moves")


def test_references_accept_a_gallery_or_a_single_path():
    """The UI hands over a gallery; an API caller passing one reference passes a path. Both have to work."""
    single = app.gallery_paths("examples/subject.png")
    assert single == ["examples/subject.png"], single
    listed = app.gallery_paths(["examples/subject.png", ("examples/plate.mp4", "caption")])
    assert listed == ["examples/subject.png", "examples/plate.mp4"], listed
    assert app.gallery_paths([{"path": "examples/subject.png"}]) == ["examples/subject.png"]
    assert app.gallery_paths(None) == [] and app.gallery_paths([]) == []
    # and the order survives, because it is what numbers them in the prompt
    many = [f"examples/{n}.png" for n in "abcde"]
    assert app.gallery_paths(many) == many
    print("gallery_paths: string, list, tuples and dicts all resolve, order preserved")


def test_repaint_reuses_what_the_plan_already_encoded():
    """Planning pays the conditioner once. Repaint must not pay it again for the same prompt and references."""
    import torch

    height, width = 768, 768
    background = np.zeros((height, width, 3), np.uint8)
    stroke = np.zeros((height, width, 4), np.uint8)
    stroke[300:520, 260:500, 3] = 255
    painted = {"layers": [stroke], "background": background}
    prompt = "a capybara walks through a snowy pine forest"

    calls = []
    original_encode, original_gen = app.encode_remote, app._generate
    load_error, pipe = app.LOAD_ERROR, app.PIPE
    app.LOAD_ERROR, app.PIPE = None, object()

    def stub_generate(request, steps, seed):
        calls.append(("generate", request["num_frames"]))
        return [f for f in request["source_video"]], torch.zeros(2, 32000), 32000

    def refuse_encode(*a, **k):
        calls.append(("encode", None))
        raise AssertionError("the conditioner was called even though the plan had already encoded this")

    app._generate, app.encode_remote = stub_generate, refuse_encode
    try:
        # what plan_and_mask would have stashed: a matching signature over prompt + references + frame count
        frames, _, _, _, _ = app.read_clip(SOURCE, 0.0, 5.0)
        num_frames = app.snap_frames(frames.shape[0] / app.FPS)
        reference = "examples/subject.png"
        conditioning = {
            "signature": app.conditioning_signature(prompt, [reference], num_frames),
            "prompt_embeds": torch.zeros(1, 8, 4),
            "text_token_tags": torch.zeros(8),
            "num_frames": num_frames,
            "reference_paths": [reference],
            "synthetic": False,
        }
        path, report = app.generate(
            SOURCE, prompt, [reference], None, 7, painted=painted, steps=4, conditioning=conditioning,
        )
        assert ("encode", None) not in calls, calls
        assert any(kind == "generate" for kind, _ in calls), calls
        print(f"reuse: conditioner skipped, {len(calls)} call(s) made -> {[k for k, _ in calls]}")

        # and a changed prompt must invalidate it
        stale = dict(conditioning, signature="not-the-same")
        try:
            app.generate(SOURCE, prompt, [reference], None, 7, painted=painted, steps=4, conditioning=stale)
        except Exception as error:
            # generate() wraps a conditioner failure in gr.Error, so the assertion arrives as its message
            assert "conditioner was called" in str(error), error
            print("invalidation: a mismatched signature re-encodes, as it must")
        else:
            raise AssertionError("a stale signature should have forced a re-encode")
    finally:
        app._generate, app.encode_remote = original_gen, original_encode
        app.LOAD_ERROR, app.PIPE = load_error, pipe


def test_the_pose_default_agrees_with_itself():
    """The checkbox and `generate`'s parameter have to carry the same default.

    They are two separate declarations of one decision: the component drives the UI, the parameter drives any caller
    that omits the argument. If they drift, a motion clip means different things through the two doors and nothing
    raises. Pose is the default because a motion clip is normally supplied in order to *change* the subject, and the
    raw footage re-supplies the old subject's appearance.
    """
    import ast
    import inspect
    import re

    source = inspect.getsource(app)
    box = re.search(r'as_pose = gr\.Checkbox\((.*?)\n                    \)', source, re.DOTALL).group(1)
    component_default = "value=True" in box
    generate = next(
        n for n in ast.walk(ast.parse(source))
        if isinstance(n, ast.FunctionDef) and n.name == "generate"
    )
    index = [a.arg for a in generate.args.args].index("as_pose")
    offset = index - (len(generate.args.args) - len(generate.args.defaults))
    parameter_default = generate.args.defaults[offset].value

    assert component_default is parameter_default is True, (component_default, parameter_default)
    print(f"pose default: checkbox {component_default}, parameter {parameter_default}")


def test_examples_can_actually_cache():
    """`gr.Examples` needs `fn` and `outputs` before `cache_examples` means anything.

    Without them the rows are input presets, gradio has nothing to store, and every click re-runs the whole thing.
    Caching also has to be `lazy`: eager caching runs every row at build time, outside any request, where there is no
    ZeroGPU token to forward. And `run_example`'s parameters are mapped onto the `inputs` list positionally, so this
    pins that order the same way `request_inputs` is pinned.
    """
    import ast
    import inspect
    import re

    source = inspect.getsource(app)
    call = re.search(r"gr\.Examples\((.*?)\n    \)", source, re.DOTALL).group(1)
    assert "fn=run_example" in call, "no fn, so nothing can be cached"
    assert "cache_examples=True" in call and 'cache_mode="lazy"' in call, call[-400:]
    assert "outputs=[" in call, "no outputs, so nothing can be cached"

    listed = re.search(r"inputs=\[(.*?)\]", call, re.DOTALL).group(1)
    names = [t.strip() for t in listed.replace("\n", " ").split(",") if t.strip()]
    runner = next(
        n for n in ast.walk(ast.parse(source))
        if isinstance(n, ast.FunctionDef) and n.name == "run_example"
    )
    params = [a.arg for a in runner.args.args if a.arg != "progress"]
    alias = {"source": "source_path", "references": "references", "video": "video_path"}
    mapped = [alias.get(name, name) for name in names]
    assert mapped == params, list(zip(mapped, params))

    # Every declared output has to be a component that exists by the time the block is built.
    outputs = [t.strip() for t in re.search(r"outputs=\[(.*?)\]", call, re.DOTALL).group(1).split(",") if t.strip()]
    returned = next(n for n in ast.walk(runner) if isinstance(n, ast.Return))
    assert len(returned.value.elts) == len(outputs), (len(returned.value.elts), outputs)
    print(f"examples: lazy cache over {len(names)} inputs -> {len(outputs)} outputs")


def test_simple_mode_offers_one_button():
    """Simple has a single button, and `toggle_mode` has to drive exactly the components it is wired to.

    Both halves are structural, so they are read off the source rather than a running Blocks: the number of updates
    `toggle_mode` returns must match the length of its `outputs` list, or gradio assigns them to the wrong
    components silently, and `plan` must be born hidden because Simple is the default view.
    """
    import ast
    import inspect
    import re

    source = inspect.getsource(app)
    tree = ast.parse(source)

    # Generate mask starts hidden: the default view is Simple, where Inpaint is the only button.
    default_view = re.search(r'view_mode = gr\.Radio\(\s*\[(.*?)\],\s*value="(\w+)"', source, re.DOTALL)
    assert default_view and default_view.group(2) == "Simple", default_view and default_view.group(2)
    button = re.search(r'plan = gr\.Button\("Generate mask".*?\)', source, re.DOTALL).group(0)
    assert "visible=False" in button, button

    toggle = next(
        n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "toggle_mode"
    )
    returned = next(n for n in ast.walk(toggle) if isinstance(n, ast.Return))
    assert isinstance(returned.value, ast.Tuple), ast.dump(returned)
    updates = len(returned.value.elts)

    wiring = re.search(r"view_mode\.change\(\s*toggle_mode,\s*\[.*?\],\s*\[(.*?)\]", source, re.DOTALL)
    outputs = [t.strip() for t in wiring.group(1).replace("\n", " ").split(",") if t.strip()]
    assert updates == len(outputs), (updates, outputs)
    assert "plan" in outputs, outputs
    print(f"simple mode: {updates} updates onto {outputs}")


def test_the_inputs_list_lines_up_with_the_signature():
    """Gradio maps an inputs list onto a function *positionally*.

    So a parameter moving, or a component leaving the list, silently shifts every argument after it rather than
    raising — which is how a checkbox once arrived where a brush was expected. This pins the two together.
    """
    import ast
    import inspect
    import re

    source = inspect.getsource(app)
    generate = next(
        n for n in ast.walk(ast.parse(source))
        if isinstance(n, ast.FunctionDef) and n.name == "generate"
    )
    params = [a.arg for a in generate.args.args if a.arg != "progress"]
    listed = re.search(r"request_inputs = \[(.*?)\]", source, re.DOTALL).group(1)
    names = [t.strip() for t in listed.replace("\n", " ").split(",") if t.strip()]
    alias = {"references": "image_1", "audio": "audio_path", "video": "video_path", "source": "source_path"}
    mapped = [alias.get(name, name) for name in names]

    assert len(mapped) <= len(params), (len(mapped), len(params))
    mismatched = [(i, a, b) for i, (a, b) in enumerate(zip(mapped, params)) if a != b]
    assert not mismatched, mismatched
    print(f"inputs line up: {len(mapped)} components onto {len(params)} parameters, in order")


def test_inpaint_fills_in_what_it_is_missing():
    """Pressing Inpaint straight away has to work, not send the user to another button.

    The filling-in lives in `prepare_if_needed` rather than in `generate`, because segmentation cannot run under
    `track_tqdm`. These are its three branches.
    """
    calls = []
    original_plan, original_pi = app.plan_and_mask, app.plan_instruction

    def fake_plan(source_path, instruction, start, duration, references, grow, cache=None, progress=None):
        calls.append("plan_and_mask")
        return ("a capybara walks through a snowy pine forest", "examples/plate_mask.mp4", "base", "view", "",
                {"x": 1}, {"signature": "s"})

    def fake_prompt_only(source_path, instruction, start, duration, references, progress=None):
        calls.append("plan_instruction")
        return "fox", "a capybara walks through a snowy pine forest"

    app.plan_and_mask, app.plan_instruction = fake_plan, fake_prompt_only
    try:
        # nothing to go on but an instruction: plan and segment
        out = app.prepare_if_needed(SOURCE, "", [], None, "replace the fox with a capybara", 0, 5, 48, None, None, None)
        assert calls == ["plan_and_mask"], calls
        assert out[0]["value"].startswith("a capybara") and out[0]["visible"] is True, out[0]
        assert out[1] == "examples/plate_mask.mp4" and out[5] == {"x": 1}
        assert len(out) == 7 and out[6] == {"signature": "s"}, out[6]
        print(f"no prompt, no mask -> {calls}, prompt revealed")

        # a mask already there: only the prompt half, and the mask is left alone
        calls.clear()
        out = app.prepare_if_needed(SOURCE, "", [], "examples/plate_mask.mp4", "a capybara", 0, 5, 48, None, "base", None)
        assert calls == ["plan_instruction"], calls
        assert out[0]["value"].startswith("a capybara")
        print(f"mask already there -> {calls}, no re-segmentation")

        # a prompt already there: nothing to do, and nothing overwritten
        calls.clear()
        out = app.prepare_if_needed(SOURCE, "my own wording", [], "examples/plate_mask.mp4", "a capybara", 0, 5, 48,
                                    {"cached": True}, "base", None)
        assert calls == [], calls
        assert out[5] == {"cached": True}
        # the State slots must carry values through, never a gr.update dict
        assert out[1] == "examples/plate_mask.mp4" and out[2] == "base", (out[1], out[2])
        print("prompt already there -> nothing re-planned, conditioning passed through")

        # nothing at all
        calls.clear()
        try:
            app.prepare_if_needed(SOURCE, "", [], None, "", 0, 5, 48, None, None, None)
        except Exception as error:
            assert "Say what should change" in str(error), error
            print("nothing at all -> asks for one line")
        else:
            raise AssertionError("an empty request should have been refused")
    finally:
        app.plan_and_mask, app.plan_instruction = original_plan, original_pi


def test_an_unchanged_request_is_not_planned_twice():
    """Pressing Generate mask again on the same clip, instruction and references must not re-plan or re-segment."""
    calls = []
    original_pi, original_fm = app.plan_instruction, app.find_mask
    app.plan_instruction = lambda *a, **k: (calls.append("plan"), ("fox", "a capybara in the snow"))[1]
    app.find_mask = lambda *a, **k: (calls.append("segment"), ("examples/plate_mask.mp4", "examples/plate_mask.mp4",
                                                              "view"))[1]
    try:
        first = app.plan_and_mask(SOURCE, "replace the fox with a capybara", 0, 5, ["examples/subject.png"], 24,
                                  None, progress=None)
        assert calls == ["plan", "segment"], calls
        cache = first[6]
        assert cache and cache["signature"], cache

        calls.clear()
        again = app.plan_and_mask(SOURCE, "replace the fox with a capybara", 0, 5, ["examples/subject.png"], 64,
                                  cache, progress=None)
        assert calls == [], f"it planned again: {calls}"
        assert again[0] == first[0], "the prompt should be the one already written"
        print("unchanged request -> nothing re-planned, mask only re-dilated")

        # but a changed instruction must invalidate it
        calls.clear()
        app.plan_and_mask(SOURCE, "replace the fox with a wolf", 0, 5, ["examples/subject.png"], 64, cache,
                          progress=None)
        assert calls == ["plan", "segment"], calls
        print("changed instruction -> planned again, as it must")
    finally:
        app.plan_instruction, app.find_mask = original_pi, original_fm


def test_duration_booking_is_bounded_and_scales():
    import torch

    def request(height, width, num_frames, tokens=400):
        return {
            "text_token_tags": torch.zeros(tokens),
            "references": [],
            "height": height,
            "width": width,
            "num_frames": num_frames,
        }

    small = app.get_duration(request(544, 960, 124), steps=28, seed=0)
    large = app.get_duration(request(768, 1344, 328), steps=28, seed=0)
    assert app.MIN_GPU_DURATION <= small <= large <= app.MAX_GPU_DURATION
    assert large > small, (small, large)
    print(f"get_duration: {small}s for 960x544x124, {large}s for 1344x768x328")


def test_duration_covers_every_kind_of_reference():
    """The video branch of `reference_rows` shipped broken because no test ever passed a reference at all.

    `resolve_canvas_size` takes a canvas budget as well as a multiple, and calling it with three arguments raises
    `TypeError` — inside `get_duration`, i.e. before the booking, so the request died with a bare `RuntimeError` and
    only for the one example that brings a motion clip. Every reference kind gets priced here.
    """
    import torch

    def duration(references):
        return app.get_duration(
            {
                "text_token_tags": torch.zeros(3087),
                "references": references,
                "height": 768,
                "width": 768,
                "num_frames": 124,
                "reference_edge": 768,
            },
            steps=8,
            seed=0,
        )

    bare = duration([])
    image = duration([("image", "examples/subject.png")])
    video = duration([("video", "examples/motion.mp4")])
    audio = duration([("audio", "examples/voice.wav")])
    for label, value in (("image", image), ("video", video), ("audio", audio)):
        assert app.MIN_GPU_DURATION <= value <= app.MAX_GPU_DURATION, (label, value)
        assert value >= bare, (label, value, bare)
    # A 60-frame clip at 1344x768 is worth far more conditioning rows than one still, so it must cost more.
    assert video > image, (video, image)
    print(f"get_duration: none {bare}s, image {image}s, video {video}s, audio {audio}s")


if __name__ == "__main__":
    for name, test in sorted(globals().items()):
        if name.startswith("test_"):
            print(f"\n--- {name}")
            test()
    print("\nOK")