|
1 | 1 | import folder_paths
|
2 | 2 | import comfy.audio_encoders.audio_encoders
|
3 | 3 | import comfy.utils
|
| 4 | +from typing_extensions import override |
| 5 | +from comfy_api.latest import ComfyExtension, io |
4 | 6 |
|
5 | 7 |
|
6 |
| -class AudioEncoderLoader: |
| 8 | +class AudioEncoderLoader(io.ComfyNode): |
7 | 9 | @classmethod
|
8 |
| - def INPUT_TYPES(s): |
9 |
| - return {"required": { "audio_encoder_name": (folder_paths.get_filename_list("audio_encoders"), ), |
10 |
| - }} |
11 |
| - RETURN_TYPES = ("AUDIO_ENCODER",) |
12 |
| - FUNCTION = "load_model" |
| 10 | + def define_schema(cls) -> io.Schema: |
| 11 | + return io.Schema( |
| 12 | + node_id="AudioEncoderLoader", |
| 13 | + category="loaders", |
| 14 | + inputs=[ |
| 15 | + io.Combo.Input( |
| 16 | + "audio_encoder_name", |
| 17 | + options=folder_paths.get_filename_list("audio_encoders"), |
| 18 | + ), |
| 19 | + ], |
| 20 | + outputs=[io.AudioEncoder.Output()], |
| 21 | + ) |
13 | 22 |
|
14 |
| - CATEGORY = "loaders" |
15 |
| - |
16 |
| - def load_model(self, audio_encoder_name): |
| 23 | + @classmethod |
| 24 | + def execute(cls, audio_encoder_name) -> io.NodeOutput: |
17 | 25 | audio_encoder_name = folder_paths.get_full_path_or_raise("audio_encoders", audio_encoder_name)
|
18 | 26 | sd = comfy.utils.load_torch_file(audio_encoder_name, safe_load=True)
|
19 | 27 | audio_encoder = comfy.audio_encoders.audio_encoders.load_audio_encoder_from_sd(sd)
|
20 | 28 | if audio_encoder is None:
|
21 | 29 | raise RuntimeError("ERROR: audio encoder file is invalid and does not contain a valid model.")
|
22 |
| - return (audio_encoder,) |
| 30 | + return io.NodeOutput(audio_encoder) |
23 | 31 |
|
24 | 32 |
|
25 |
| -class AudioEncoderEncode: |
| 33 | +class AudioEncoderEncode(io.ComfyNode): |
26 | 34 | @classmethod
|
27 |
| - def INPUT_TYPES(s): |
28 |
| - return {"required": { "audio_encoder": ("AUDIO_ENCODER",), |
29 |
| - "audio": ("AUDIO",), |
30 |
| - }} |
31 |
| - RETURN_TYPES = ("AUDIO_ENCODER_OUTPUT",) |
32 |
| - FUNCTION = "encode" |
33 |
| - |
34 |
| - CATEGORY = "conditioning" |
| 35 | + def define_schema(cls) -> io.Schema: |
| 36 | + return io.Schema( |
| 37 | + node_id="AudioEncoderEncode", |
| 38 | + category="conditioning", |
| 39 | + inputs=[ |
| 40 | + io.AudioEncoder.Input("audio_encoder"), |
| 41 | + io.Audio.Input("audio"), |
| 42 | + ], |
| 43 | + outputs=[io.AudioEncoderOutput.Output()], |
| 44 | + ) |
35 | 45 |
|
36 |
| - def encode(self, audio_encoder, audio): |
| 46 | + @classmethod |
| 47 | + def execute(cls, audio_encoder, audio) -> io.NodeOutput: |
37 | 48 | output = audio_encoder.encode_audio(audio["waveform"], audio["sample_rate"])
|
38 |
| - return (output,) |
| 49 | + return io.NodeOutput(output) |
| 50 | + |
| 51 | + |
| 52 | +class AudioEncoder(ComfyExtension): |
| 53 | + @override |
| 54 | + async def get_node_list(self) -> list[type[io.ComfyNode]]: |
| 55 | + return [ |
| 56 | + AudioEncoderLoader, |
| 57 | + AudioEncoderEncode, |
| 58 | + ] |
39 | 59 |
|
40 | 60 |
|
41 |
| -NODE_CLASS_MAPPINGS = { |
42 |
| - "AudioEncoderLoader": AudioEncoderLoader, |
43 |
| - "AudioEncoderEncode": AudioEncoderEncode, |
44 |
| -} |
| 61 | +async def comfy_entrypoint() -> AudioEncoder: |
| 62 | + return AudioEncoder() |
0 commit comments