199 lines
11 KiB
Python
199 lines
11 KiB
Python
"""Real ComfyUI V3 bridge mapping, file ownership and exact CPU round trips."""
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def test_real_media_bridge_nodes():
|
|
core = os.environ.get("COMFYUI_SOURCE_ROOT")
|
|
if core is None:
|
|
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU media bridge acceptance")
|
|
assert core and (Path(core) / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
|
|
"Configured COMFYUI_SOURCE_ROOT must provide native video classes"
|
|
)
|
|
probe = r'''
|
|
import asyncio, importlib.util, io, json, os, sys, tempfile
|
|
from fractions import Fraction
|
|
import av
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
import torch
|
|
sys.path.insert(0, sys.argv[1])
|
|
import comfy.cli_args
|
|
comfy.cli_args.args.cpu = True
|
|
comfy.cli_args.args.disable_assets = True
|
|
def forbidden(*args, **kwargs):
|
|
raise AssertionError('Bridge tests must not query or initialize a GPU')
|
|
torch.cuda.is_available = lambda: False
|
|
torch.cuda.device_count = lambda: 0
|
|
for name in ('init', '_lazy_init', 'current_device'):
|
|
setattr(torch.cuda, name, forbidden)
|
|
import execution
|
|
import folder_paths
|
|
from comfy_api.latest._input_impl.video_types import VideoFromComponents, VideoFromFile, VideoFromList
|
|
from comfy_api.latest._util import VideoComponents
|
|
for name in ('is_available', 'device_count'):
|
|
setattr(torch.cuda, name, forbidden)
|
|
spec = importlib.util.spec_from_file_location('typed_bridge_probe', Path(sys.argv[2])/'serverless_bridge/__init__.py', submodule_search_locations=[str(Path(sys.argv[2])/'serverless_bridge')])
|
|
bridge = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = bridge
|
|
spec.loader.exec_module(bridge)
|
|
from typed_bridge_probe.nodes import MediaImportV1, MediaExportV1, export_values, import_value
|
|
from typed_bridge_probe.contract import SOCKET_KINDS
|
|
from typed_bridge_probe.video_codec import encode_video
|
|
|
|
def tiny_video():
|
|
buffer = io.BytesIO()
|
|
with av.open(buffer, 'w', format='matroska') as container:
|
|
stream = container.add_stream('ffv1', rate=Fraction(24, 1))
|
|
stream.width = stream.height = 32
|
|
stream.pix_fmt = 'yuv420p'
|
|
for level in (16, 64, 128):
|
|
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), level, dtype=np.uint8), format='rgb24')
|
|
for packet in stream.encode(frame): container.mux(packet)
|
|
for packet in stream.encode(): container.mux(packet)
|
|
return buffer.getvalue()
|
|
|
|
def exact_tensor(actual, expected):
|
|
assert actual.device.type == 'cpu' and actual.dtype == expected.dtype
|
|
assert actual.shape == expected.shape
|
|
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
|
|
|
def exact_video(actual, expected, wire):
|
|
assert type(actual) is type(expected)
|
|
assert encode_video(actual) == wire
|
|
if type(actual) is VideoFromComponents:
|
|
a = vars(actual)['_VideoFromComponents__components']
|
|
e = vars(expected)['_VideoFromComponents__components']
|
|
assert type(a) is VideoComponents and type(a.frame_rate) is Fraction
|
|
assert a.frame_rate == e.frame_rate and a.metadata == e.metadata
|
|
assert a.metadata is not e.metadata
|
|
for name in ('images', 'alpha'): exact_tensor(getattr(a, name), getattr(e, name))
|
|
exact_tensor(a.audio['waveform'], e.audio['waveform'])
|
|
assert a.audio['sample_rate'] == e.audio['sample_rate']
|
|
assert actual.get_bit_depth() == 10 and actual.get_color_space() == 'HDR PQ'
|
|
else:
|
|
values = actual.videos if type(actual) is VideoFromList else [actual]
|
|
for value in values:
|
|
buffer = vars(value)['_VideoFromFile__file']
|
|
assert type(buffer) is io.BytesIO and not buffer.closed
|
|
buffer.seek(0)
|
|
with av.open(buffer) as container:
|
|
assert len(list(container.decode(video=0))) == 3
|
|
if type(actual) is VideoFromList:
|
|
exact_tensor(actual.complete_audio['waveform'], expected.complete_audio['waveform'])
|
|
|
|
async def check():
|
|
extension = await bridge.comfy_entrypoint()
|
|
assert await extension.get_node_list() == [MediaImportV1, MediaExportV1]
|
|
imp = MediaImportV1.GET_NODE_INFO_V1()
|
|
exp = MediaExportV1.GET_NODE_INFO_V1()
|
|
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO','*'], imp
|
|
assert exp['output_node'] and MediaExportV1.define_schema().is_input_list, exp
|
|
with tempfile.TemporaryDirectory(prefix='typed-media-', dir=os.environ.get('TMPDIR')) as scratch:
|
|
base = Path(scratch)
|
|
roots = {key: base/key for key in ('input','output','temp')}
|
|
for path in roots.values(): path.mkdir()
|
|
def annotated(name):
|
|
for key, root in roots.items():
|
|
suffix = f' [{key}]'
|
|
if name.endswith(suffix): return str(root/name[:-len(suffix)])
|
|
return str(roots['input']/name)
|
|
with patch.object(folder_paths,'get_input_directory',lambda: str(roots['input'])), patch.object(folder_paths,'get_output_directory',lambda: str(roots['output'])), patch.object(folder_paths,'get_temp_directory',lambda: str(roots['temp'])), patch.object(folder_paths,'get_annotated_filepath',annotated):
|
|
examples = {'images': torch.rand(2,2,3,4,dtype=torch.float64),
|
|
'masks': torch.tensor([[[0.123456789,0.987654321]]],dtype=torch.float64),
|
|
'audio': {'waveform':torch.randn(1,6,31,dtype=torch.float64),'sample_rate':44100},
|
|
'texts': ' café\n🎥\r\n'}
|
|
media = tiny_video()
|
|
sources = [roots['input']/f'source-{n}.mkv' for n in range(3)]
|
|
for path in sources: path.write_bytes(media)
|
|
images = torch.arange(36, dtype=torch.float64, device='cpu').reshape(2,2,3,3) / 997
|
|
images[0,0,0,0] = -0.0
|
|
component_video = VideoFromComponents(VideoComponents(
|
|
images=images, frame_rate=Fraction(30000,1001),
|
|
alpha=torch.full((2,2,3,1),0.123456789,dtype=torch.float64,device='cpu'),
|
|
audio=examples['audio'], metadata={'title':'雪', 'nested':[True,None,-4,0.125]}),
|
|
bit_depth=10,color_space='HDR PQ')
|
|
videos = [component_video,
|
|
VideoFromFile(str(sources[0]),start_time=-0.0,duration=0,crop=(0,0,2,2)),
|
|
VideoFromList([VideoFromFile(str(sources[1]),start_time=-1.25,duration=-2.5,crop=(-3,5,7,-9)),
|
|
VideoFromFile(str(sources[2]),start_time=0,duration=0)],
|
|
complete_audio=examples['audio'])]
|
|
for socket,value in [*examples.items(), *[('video',v) for v in videos]]:
|
|
kind = SOCKET_KINDS[socket]
|
|
expected_wire = encode_video(value) if kind == 'VIDEO' else None
|
|
binding = dict(binding_id='b-'+socket,collector_id='collector',job_id='job',participant_id='worker',node_id='export-'+socket,socket=socket,kind=kind)
|
|
controls = {'request_id':'request','binding_json':json.dumps(binding)}
|
|
data, missing, v3 = execution.get_input_data(controls,MediaExportV1,'export-'+socket)
|
|
assert not missing, missing
|
|
data[socket] = [value]
|
|
results = await execution._async_map_node_over_list('probe','export-'+socket,MediaExportV1,data,MediaExportV1.FUNCTION,v3_data=v3)
|
|
results = await execution.resolve_map_node_over_list_results(results)
|
|
raw = results[0].ui
|
|
exported = raw['media_exports'][0]
|
|
assert exported['binding']['count'] == 1
|
|
artifact = exported['artifacts'][0]
|
|
transport = artifact.pop('transport')
|
|
filename = transport['subfolder']+'/'+transport['filename']+' [temp]'
|
|
decoded = import_value(artifact,filename)
|
|
if kind in ('IMAGE','MASK'):
|
|
assert decoded.dtype == value.dtype and torch.equal(decoded,value)
|
|
elif kind == 'AUDIO':
|
|
assert decoded['sample_rate'] == value['sample_rate']
|
|
assert torch.equal(decoded['waveform'],value['waveform'])
|
|
elif kind == 'VIDEO': exact_video(decoded,value,expected_wire)
|
|
else: assert decoded == value
|
|
mapped, missing, hidden = execution.get_input_data({'artifact_json':json.dumps(artifact),'filename':filename},MediaImportV1,'import')
|
|
got = await execution._async_map_node_over_list('probe','import',MediaImportV1,mapped,MediaImportV1.FUNCTION,v3_data=hidden)
|
|
got = await execution.resolve_map_node_over_list_results(got)
|
|
assert len(got[0].result) == 6
|
|
assert all(x is None for n,x in enumerate(got[0].result) if n != list(SOCKET_KINDS).index(socket))
|
|
bad = dict(artifact,sha256='0'*64)
|
|
try: import_value(bad,filename)
|
|
except ValueError: pass
|
|
else: raise AssertionError('checksum accepted')
|
|
try: import_value(artifact,str(base/'outside.bin'))
|
|
except ValueError: pass
|
|
else: raise AssertionError('outside path accepted')
|
|
if kind == 'VIDEO':
|
|
mapped_video = got[0].result[list(SOCKET_KINDS).index(socket)]
|
|
# Remove both the exported artifact and the original input
|
|
# containers before checking independently owned BytesIOs.
|
|
Path(annotated(filename)).unlink()
|
|
originals = value.videos if type(value) is VideoFromList else [value]
|
|
for original in originals:
|
|
if type(original) is VideoFromFile:
|
|
Path(vars(original)['_VideoFromFile__file']).unlink()
|
|
exact_video(decoded,value,expected_wire)
|
|
exact_video(mapped_video,value,expected_wire)
|
|
if type(value) is VideoFromFile:
|
|
assert tuple(decoded.get_components().images.shape) == (3,2,2,3)
|
|
text_binding = dict(binding_id='empty',collector_id='collector',job_id='job',participant_id='worker',node_id='export-empty',socket='texts',kind='STRING')
|
|
out = export_values('request',text_binding,[''])
|
|
assert out['artifacts'][0]['byte_length'] == 0
|
|
before = set(roots['temp'].rglob('*.bin'))
|
|
try: export_values('request',text_binding,['first',None])
|
|
except ValueError: pass
|
|
else: raise AssertionError('invalid text accepted')
|
|
assert set(roots['temp'].rglob('*.bin')) == before
|
|
print('MEDIA_BRIDGE_CPU_OK')
|
|
asyncio.run(check())
|
|
'''
|
|
env = dict(os.environ)
|
|
env.pop("PYTHONPATH", None)
|
|
env.pop("COMFYUI_SOURCE_ROOT", None)
|
|
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
|
scratch.mkdir(parents=True, exist_ok=True)
|
|
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", PYTHONDONTWRITEBYTECODE="1",
|
|
TMPDIR=str(scratch))
|
|
result = subprocess.run([sys.executable, "-c", probe, str(Path(core).resolve()), str(ROOT)],
|
|
capture_output=True, text=True, timeout=90, env=env)
|
|
assert result.returncode == 0, result.stdout + result.stderr
|
|
assert "MEDIA_BRIDGE_CPU_OK" in result.stdout
|