mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-20 22:56:41 +08:00
61 lines
2.8 KiB
Python
61 lines
2.8 KiB
Python
import subprocess
|
|
import json
|
|
import os
|
|
import torch
|
|
import shutil
|
|
|
|
import server
|
|
import folder_paths
|
|
|
|
web = server.web
|
|
|
|
@server.PromptServer.instance.routes.post("/VHS_test")
|
|
async def test(request):
|
|
try:
|
|
req_data = await request.json()
|
|
output = req_data['output']['gifs'][0]
|
|
filename = output['filename']
|
|
typ = output['type']
|
|
base_args = ["ffprobe", "-v", "error", '-count_packets', "-show_entries", "stream", "-of", "json"]
|
|
video = folder_paths.get_annotated_filepath(f'{filename} [{typ}]')
|
|
vprobe = json.loads(subprocess.run(base_args + ['-select_streams', 'v:0', video],
|
|
capture_output=True, check=True).stdout)['streams'][0]
|
|
aprobe = json.loads(subprocess.run(base_args + ['-select_streams', 'a:0', video],
|
|
capture_output=True, check=True).stdout)['streams']
|
|
probe = {'video': vprobe}
|
|
if len(aprobe) > 0:
|
|
probe['audio'] = aprobe[0]
|
|
errors = []
|
|
compare = None
|
|
for test in req_data['tests']:
|
|
if test['type'] == 'compare':
|
|
compare = test
|
|
continue
|
|
key = test['key']
|
|
expected = test['value']
|
|
actual = probe[test['type']][key]
|
|
if expected != actual:
|
|
#Consider always dumping type?
|
|
errors.append(f'{key}: {expected} != {actual}')
|
|
if len(errors) == 0 and compare is not None:
|
|
if not os.path.exists(compare['filename']):
|
|
os.makedirs(os.path.split(compare['filename'])[0], exist_ok=True)
|
|
shutil.copy(video, compare['filename'])
|
|
print("Missing comparison file has been initialized from output:", os.path.abspath(compare['filename']))
|
|
else:
|
|
#NOTE: This does not include the full memory optimizations of VHS
|
|
#Tests should be small
|
|
#TODO: Figure out way to do opacity comparison. May need to do blending in python
|
|
#(easy, but slower and more memory intensive)
|
|
diff = subprocess.run(['ffmpeg', '-v', 'error', '-i', video, '-i', compare['filename'], '-filter_complex', 'blend=all_mode=grainextract', '-pix_fmt', 'rgb24', '-f', 'rawvideo', '-'], stdout=subprocess.PIPE, check=True).stdout
|
|
diff = torch.frombuffer(diff, dtype=torch.uint8).to(dtype=torch.float32).div_(255)
|
|
#diff = diff.reshape((-1,4))
|
|
d = (diff-0.5).abs().sum()/diff.size(0)
|
|
if d > compare['tolerance']:
|
|
errors.append(f'Similarity is outside specified tolerance: {d}')
|
|
else:
|
|
print('d:', d)
|
|
return web.json_response(errors)
|
|
except Exception as e:
|
|
return web.json_response(str(e))
|