diff --git a/__init__.py b/__init__.py index d2f1739..73b1be5 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,5 @@ # ComfyUI-DiaTTS/__init__.py -from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .nodes import comfy_entrypoint -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] +__all__ = ['comfy_entrypoint'] diff --git a/nodes.py b/nodes.py index 52cf5c2..f607573 100644 --- a/nodes.py +++ b/nodes.py @@ -10,6 +10,8 @@ from safetensors.torch import load_file as load_safetensors_file import comfy.utils import torchaudio # Needed for potential resampling in encode_audio_prompt +from comfy_api.latest import ComfyExtension, io, ui + # --- Import Dia library components --- try: from .dia_lib.model import Dia, ComputeDtype, DEFAULT_SAMPLE_RATE @@ -41,29 +43,35 @@ def get_torch_device(): # --- Global model cache --- loaded_dia_objects = {} -class DiaLoader: +class DiaLoader(io.ComfyNode): """Loads the Dia-1.6B TTS model from a local safetensors file.""" + @classmethod - def INPUT_TYPES(s): + def define_schema(cls) -> io.Schema: """Finds .safetensors files in diffusion_models directories.""" try: safetensors_files = folder_paths.get_filename_list("diffusion_models") - s.dia_model_files = sorted([f for f in safetensors_files if f.lower().endswith(".safetensors")]) - if not s.dia_model_files: - print("DiaLoader: No .safetensors files found in diffusion_models directories.") - s.dia_model_files = ["None"] + dia_model_files = sorted([f for f in safetensors_files if f.lower().endswith(".safetensors")]) + if not dia_model_files: + dia_model_files = ["None"] except Exception as e: print(f"DiaLoader: Warning - Could not access diffusion_models paths: {e}") - s.dia_model_files = ["None"] + dia_model_files = ["None"] - return { "required": { "ckpt_name": (s.dia_model_files,), }, } + return io.Schema( + node_id="DiaLoader", + display_name="Dia 1.6b Loader", + category="audio/DiaTTS", + inputs=[ + io.Combo.Input("ckpt_name", options=dia_model_files) + ], + outputs=[ + io.Custom("DIA_MODEL").Output(display_name="dia_model") + ] + ) - RETURN_TYPES = ("DIA_MODEL",) - RETURN_NAMES = ("dia_model",) - FUNCTION = "load_dia_model" - CATEGORY = "audio/DiaTTS" - - def load_dia_model(self, ckpt_name: str): + @classmethod + def execute(cls, ckpt_name: str) -> io.NodeOutput: """Loads the safetensors weights, combines with embedded config, loads DAC, and prepares the Dia object.""" global loaded_dia_objects @@ -154,41 +162,38 @@ class DiaLoader: print(f"DiaLoader: Caching loaded Dia object.") loaded_dia_objects[current_key] = dia_object - return (dia_object,) + return io.NodeOutput(dia_object) -class DiaGenerate: +class DiaGenerate(io.ComfyNode): """Generates audio using a pre-loaded Dia TTS model, optionally with an audio prompt.""" + @classmethod - def INPUT_TYPES(s): - """Defines the inputs for audio generation.""" - return { - "required": { - "dia_model": ("DIA_MODEL",), - "text": ("STRING", { - "multiline": True, - "dynamicPrompts": False, - "default": "" - }), - "max_tokens": ("INT", {"default": 1720, "min": 860, "max": 3072, "step": 10}), - "cfg_scale": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 7.0, "step": 0.1}), - "temperature": ("FLOAT", {"default": 1.3, "min": 0.1, "max": 1.5, "step": 0.05}), - "top_p": ("FLOAT", {"default": 0.95, "min": 0.1, "max": 1.0, "step": 0.01}), - "cfg_filter_top_k": ("INT", {"default": 32, "min": 1, "max": 100, "step": 1}), - "speed_factor": ("FLOAT", {"default": 0.94, "min": 0.5, "max": 1.5, "step": 0.01}), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - }, - "optional": { - "audio_prompt": ("AUDIO",), # Optional audio prompt input - } - } + def define_schema(cls) -> io.Schema: + """Defines the inputs and outputs for audio generation.""" + return io.Schema( + node_id="DiaGenerate", + display_name="Dia TTS Generate", + category="audio/DiaTTS", + inputs=[ + io.Custom("DIA_MODEL").Input("dia_model"), + io.String.Input("text", multiline=True, default=""), + io.Int.Input("max_tokens", default=1720, min=860, max=3072, step=10), + io.Float.Input("cfg_scale", default=3.0, min=1.0, max=7.0, step=0.1), + io.Float.Input("temperature", default=1.3, min=0.1, max=1.5, step=0.05), + io.Float.Input("top_p", default=0.95, min=0.1, max=1.0, step=0.01), + io.Int.Input("cfg_filter_top_k", default=32, min=1, max=100, step=1), + io.Float.Input("speed_factor", default=0.94, min=0.5, max=1.5, step=0.01), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.Custom("AUDIO").Input("audio_prompt", optional=True) + ], + outputs=[ + io.Custom("AUDIO").Output(display_name="audio") + ] + ) - RETURN_TYPES = ("AUDIO",) - RETURN_NAMES = ("audio",) - FUNCTION = "generate_audio" - CATEGORY = "audio/DiaTTS" - - def generate_audio(self, dia_model: Dia, text: str, max_tokens: int, cfg_scale: float, temperature: float, top_p: float, cfg_filter_top_k: int, speed_factor: float, seed: int, audio_prompt=None): + @classmethod + def execute(cls, dia_model: Dia, text: str, max_tokens: int, cfg_scale: float, temperature: float, top_p: float, cfg_filter_top_k: int, speed_factor: float, seed: int, audio_prompt=None) -> io.NodeOutput: """ Performs TTS generation. If audio_prompt is provided, the 'text' input should contain the transcript of the audio_prompt followed by the text to generate. @@ -277,7 +282,7 @@ class DiaGenerate: if output_np is None or output_np.size == 0: print("DiaGenerate: Warning - Generation returned empty. Outputting silence.") silent_tensor = torch.zeros((1, 1, DEFAULT_SAMPLE_RATE), dtype=torch.float32) # Shape [1, 1, T] - return ({'waveform': silent_tensor, 'sample_rate': DEFAULT_SAMPLE_RATE},) + return io.NodeOutput({'waveform': silent_tensor, 'sample_rate': DEFAULT_SAMPLE_RATE}) # --- Speed Factor Adjustment --- if speed_factor != 1.0: @@ -320,7 +325,7 @@ class DiaGenerate: raise ValueError(f"Final tensor has a zero dimension: {output_tensor.shape}") result = {'waveform': output_tensor, 'sample_rate': DEFAULT_SAMPLE_RATE} - return (result,) + return io.NodeOutput(result) except Exception as format_e: print(f"DiaGenerate: Error formatting output: {format_e}") traceback.print_exc() @@ -337,13 +342,9 @@ class DiaGenerate: torch.cuda.empty_cache() -# --- Node Mappings --- -NODE_CLASS_MAPPINGS = { - "DiaLoader": DiaLoader, - "DiaGenerate": DiaGenerate, -} +class DiaExtension(ComfyExtension): + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [DiaLoader, DiaGenerate] -NODE_DISPLAY_NAME_MAPPINGS = { - "DiaLoader": "Dia 1.6b Loader", - "DiaGenerate": "Dia TTS Generate", -} +async def comfy_entrypoint() -> DiaExtension: + return DiaExtension()