89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
import wave
|
|
|
|
import torch
|
|
import torch.nn.functional as torch_functional
|
|
|
|
|
|
def install_comfy_stubs() -> None:
|
|
"""Install the narrow ComfyUI seams exercised by the reframe unit tests."""
|
|
|
|
folder_paths = ModuleType("folder_paths")
|
|
folder_paths.get_input_directory = lambda: str(Path.cwd())
|
|
folder_paths.annotated_filepath = lambda name: (name, None)
|
|
folder_paths.filter_files_content_types = lambda files, _types: list(files)
|
|
|
|
def load_audio(path: str):
|
|
with wave.open(str(path), "rb") as handle:
|
|
if handle.getsampwidth() != 2:
|
|
raise ValueError("test audio stub supports 16-bit PCM only")
|
|
channels = handle.getnchannels()
|
|
frame_count = handle.getnframes()
|
|
sample_rate = handle.getframerate()
|
|
payload = bytearray(handle.readframes(frame_count))
|
|
samples = torch.frombuffer(payload, dtype=torch.int16).reshape(-1, channels)
|
|
waveform = samples.transpose(0, 1).to(torch.float32).div_(32768.0).contiguous()
|
|
return waveform, sample_rate
|
|
|
|
comfy_extras = ModuleType("comfy_extras")
|
|
comfy_extras.__path__ = []
|
|
comfy_audio = ModuleType("comfy_extras.nodes_audio")
|
|
comfy_audio.load = load_audio
|
|
comfy_extras.nodes_audio = comfy_audio
|
|
|
|
comfy = ModuleType("comfy")
|
|
comfy.__path__ = []
|
|
comfy_utils = ModuleType("comfy.utils")
|
|
comfy_utils.common_upscale = lambda images, width, height, _method, _crop: (
|
|
torch_functional.interpolate(
|
|
images,
|
|
size=(int(height), int(width)),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
)
|
|
comfy.utils = comfy_utils
|
|
|
|
class UnsupportedVideo:
|
|
def __init__(self, *_args, **_kwargs):
|
|
raise RuntimeError("video decoding is outside this unit-test seam")
|
|
|
|
comfy_latest = ModuleType("comfy_api.latest")
|
|
comfy_latest.InputImpl = SimpleNamespace(
|
|
VideoFromFile=UnsupportedVideo,
|
|
VideoFromComponents=UnsupportedVideo,
|
|
)
|
|
comfy_latest.Types = SimpleNamespace(
|
|
VideoComponents=lambda **values: SimpleNamespace(**values)
|
|
)
|
|
comfy_api = ModuleType("comfy_api")
|
|
comfy_api.__path__ = []
|
|
comfy_api.latest = comfy_latest
|
|
|
|
aiohttp_web = ModuleType("aiohttp.web")
|
|
aiohttp_web.json_response = lambda data, status=200: SimpleNamespace(
|
|
data=data,
|
|
status=status,
|
|
)
|
|
aiohttp = ModuleType("aiohttp")
|
|
aiohttp.__path__ = []
|
|
aiohttp.web = aiohttp_web
|
|
|
|
sys.modules.update(
|
|
{
|
|
"aiohttp": aiohttp,
|
|
"aiohttp.web": aiohttp_web,
|
|
"comfy": comfy,
|
|
"comfy.utils": comfy_utils,
|
|
"comfy_api": comfy_api,
|
|
"comfy_api.latest": comfy_latest,
|
|
"comfy_extras": comfy_extras,
|
|
"comfy_extras.nodes_audio": comfy_audio,
|
|
"folder_paths": folder_paths,
|
|
}
|
|
)
|