ComfyUI/tests-unit/comfy_test/test_vae_decode_tiled_nested.py

from unittest.mock import MagicMock

import torch

from comfy.cli_args import args as cli_args

if not torch.cuda.is_available():
    cli_args.cpu = True

import comfy.nested_tensor  # noqa: E402
import nodes  # noqa: E402


def test_vae_decode_tiled_unwraps_nested_tensor():
    video = torch.zeros(1, 4, 2, 8, 8)
    audio = torch.zeros(1, 2, 2, 40)
    samples = {"samples": comfy.nested_tensor.NestedTensor((video, audio))}

    vae = MagicMock()
    vae.temporal_compression_decode.return_value = None
    vae.spacial_compression_decode.return_value = 8
    vae.decode_tiled.return_value = torch.zeros(1, 3, 2, 8, 8)

    nodes.VAEDecodeTiled().decode(vae, samples, tile_size=512)

    decoded_arg = vae.decode_tiled.call_args[0][0]
    assert decoded_arg is video