import importlib.util import sys import types from pathlib import Path class _Port: def __init__(self, name=None, **kwargs): self.name = name self.kwargs = kwargs class _PortType: @staticmethod def Input(name, **kwargs): return _Port(name, **kwargs) @staticmethod def Output(name=None, **kwargs): return _Port(name, **kwargs) class _NodeOutput: def __init__(self, *values): self.values = values class _Schema: def __init__(self, **kwargs): self.__dict__.update(kwargs) def _load_node_module(): io = types.SimpleNamespace( Audio=_PortType, Boolean=_PortType, Combo=_PortType, ComfyNode=object, Int=_PortType, NodeOutput=_NodeOutput, Schema=_Schema, String=_PortType, Video=_PortType, ) comfy_api = types.ModuleType("comfy_api") comfy_api_latest = types.ModuleType("comfy_api.latest") comfy_api_latest.io = io comfy_api_latest.Types = types.SimpleNamespace( VideoContainer=types.SimpleNamespace(AUTO="auto"), VideoCodec=types.SimpleNamespace(AUTO="auto"), ) comfy_api.latest = comfy_api_latest package = types.ModuleType("easy_media") package.__path__ = [] nodes_package = types.ModuleType("easy_media.nodes") nodes_package.__path__ = [] modules_package = types.ModuleType("easy_media.modules") modules_package.__path__ = [] asr_package = types.ModuleType("easy_media.modules.asr") asr_package.__path__ = [] recognition_module = types.ModuleType("easy_media.modules.asr.subtitle_recognition") recognition_module.SUBTITLE_RECOGNITION_METHODS = ["qwen3-asr", "whisper-large-v3"] recognition_module.recognize_audio_subtitles = lambda *args: [] utils_module = types.ModuleType("easy_media.utils") utils_module.extract_video_audio_to_temp = lambda *args, **kwargs: None utils_module.save_audio_to_temp_wav = lambda *args, **kwargs: None utils_module.subtitle_segments_to_srt = lambda segments: "" utils_module.subtitle_segments_to_timestamp_text = lambda segments: "" utils_module.video_input_to_local_file = lambda *args, **kwargs: ("", []) sys.modules.update({ "comfy_api": comfy_api, "comfy_api.latest": comfy_api_latest, "easy_media": package, "easy_media.nodes": nodes_package, "easy_media.modules": modules_package, "easy_media.modules.asr": asr_package, "easy_media.modules.asr.subtitle_recognition": recognition_module, "easy_media.utils": utils_module, }) path = Path(__file__).parents[1] / "nodes" / "subtitle.py" spec = importlib.util.spec_from_file_location("easy_media.nodes.subtitle", path) module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) return module def test_recognize_subtitle_schema_supports_two_models_and_audio_or_video(): module = _load_node_module() schema = module.RecognizeSubtitle.define_schema() assert schema.node_id == "easy recognizeSubtitle" assert [item.name for item in schema.inputs] == [ "audio", "video", "model_type", "output_format", "max_sentence_length", "unload_model", ] assert schema.inputs[0].kwargs["optional"] is True assert schema.inputs[1].kwargs["optional"] is True assert schema.inputs[2].kwargs["options"] == ["qwen3-asr", "whisper-large-v3"] assert schema.inputs[3].kwargs["options"] == ["srt", "timestamp"] assert schema.inputs[4].kwargs["default"] == 20 assert schema.inputs[5].kwargs["default"] is True assert [item.name for item in schema.outputs] == ["SUBTITLE_TEXT"] def test_recognize_subtitle_uses_audio_and_returns_srt(monkeypatch, tmp_path): module = _load_node_module() audio_path = tmp_path / "audio.wav" audio_path.write_bytes(b"wav") calls = {} monkeypatch.setattr(module, "save_audio_to_temp_wav", lambda audio: audio_path) def fake_recognize(path, model, max_sentence_length, unload_model): calls["recognize"] = (path, model, max_sentence_length, unload_model) return [{"start": 0.0, "end": 1.0, "text": "hello"}] monkeypatch.setattr(module, "recognize_audio_subtitles", fake_recognize) monkeypatch.setattr( module, "subtitle_segments_to_srt", lambda segments: "1\n00:00:00,000 --> 00:00:01,000\nhello\n", ) result = module.RecognizeSubtitle.execute( {"waveform": object(), "sample_rate": 16000}, None, "whisper-large-v3", "srt", 24, False, ) assert calls["recognize"] == (audio_path, "whisper-large-v3", 24, False) assert result.values == ("1\n00:00:00,000 --> 00:00:01,000\nhello\n",) assert not audio_path.exists() def test_recognize_subtitle_extracts_video_audio_and_cleans_temp_files(monkeypatch, tmp_path): module = _load_node_module() video_path = tmp_path / "video.mp4" copied_video_path = tmp_path / "copied.mp4" audio_path = tmp_path / "audio.wav" copied_video_path.write_bytes(b"video") audio_path.write_bytes(b"wav") monkeypatch.setattr( module, "video_input_to_local_file", lambda video, **kwargs: (str(video_path), [str(copied_video_path)]), ) monkeypatch.setattr(module, "extract_video_audio_to_temp", lambda path: audio_path) monkeypatch.setattr(module, "recognize_audio_subtitles", lambda path, model, maximum, unload: []) result = module.RecognizeSubtitle.execute(None, object(), "qwen3-asr", "srt", 48, True) assert result.values == ("",) assert not copied_video_path.exists() assert not audio_path.exists() def test_recognize_subtitle_returns_timestamp_format(monkeypatch, tmp_path): module = _load_node_module() audio_path = tmp_path / "audio.wav" audio_path.write_bytes(b"wav") monkeypatch.setattr(module, "save_audio_to_temp_wav", lambda audio: audio_path) monkeypatch.setattr( module, "recognize_audio_subtitles", lambda path, model, maximum, unload: [{"start": 0.42, "end": 1.4, "text": "hello"}], ) monkeypatch.setattr( module, "subtitle_segments_to_timestamp_text", lambda segments: "(0.42, 1.4) hello", ) result = module.RecognizeSubtitle.execute( {"waveform": object(), "sample_rate": 16000}, None, "qwen3-asr", "timestamp", 48, True, ) assert result.values == ("(0.42, 1.4) hello",)