mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-16 14:36:32 +08:00
6233790c6d
VAEDecode unwraps a NestedTensor latent (video/audio pair) to its video component before calling vae.decode(). VAEDecodeTiled skipped this unwrap and passed the NestedTensor straight into vae.decode_tiled(), which fails deep in the MiniMax H3 video VAE when a real tensor's .to() is called with the NestedTensor as an argument. Fixes #15468.
28 lines
770 B
Python
28 lines
770 B
Python
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
|