From 82f62dcbf4e48909215490232efc1d9c3cb3e344 Mon Sep 17 00:00:00 2001 From: Austin Mroz Date: Sun, 6 Apr 2025 11:15:25 -0500 Subject: [PATCH] Support audio in unbatch node --- videohelpersuite/nodes.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/videohelpersuite/nodes.py b/videohelpersuite/nodes.py index be64d4e..e2301cf 100755 --- a/videohelpersuite/nodes.py +++ b/videohelpersuite/nodes.py @@ -956,7 +956,10 @@ class Unbatch: return (torch.cat(batched),) if isinstance(batched[0], dict): out = batched[0].copy() - out['samples'] = torch.cat([x['samples'] for x in batched]) + if 'samples' in out: + out['samples'] = torch.cat([x['samples'] for x in batched]) + if 'waveform' in out: + out['waveform'] = torch.cat([x['waveform'] for x in batched]) out.pop('batch_index', None) return (out,) return (functools.reduce(lambda x,y: x+y, batched),)