diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml new file mode 100644 index 0000000..ef9d8d9 --- /dev/null +++ b/.github/workflows/publish_action.yml @@ -0,0 +1,22 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - master + - main + paths: + - "pyproject.toml" + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/CSMNode.py b/CSMNode.py new file mode 100644 index 0000000..b800eb7 --- /dev/null +++ b/CSMNode.py @@ -0,0 +1,510 @@ +from dataclasses import dataclass +from typing import List, Tuple +import ast +import silentcipher +import torch +import torchaudio +import os + +from huggingface_hub import hf_hub_download +from .models import Model, ModelArgs +from moshi.models import loaders +from tokenizers.processors import TemplateProcessing +from transformers import AutoTokenizer +import folder_paths + + +models_dir = folder_paths.models_dir + +class AddWatermark: + if torch.backends.mps.is_available(): + device = "mps" + elif torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + @classmethod + def INPUT_TYPES(s): + return {"required": { + "audio": ("AUDIO",), + "add_watermark": ("BOOLEAN", {"default": False, "tooltip": "Add watermark or not."}), + "key": ("STRING", {"default": "[212, 211, 146, 56, 201]", "tooltip": "List of integers such as [212, 211, 146, 56, 201]"}), + }, + # "optional": { + # "check_watermark": ("BOOLEAN", {"default": False, "tooltip": "Check if the audio contains watermark."}), + # } + } + + + CATEGORY = "MW_CSM" + RETURN_TYPES = ("AUDIO", "STRING") + RETURN_NAMES = ("audio", "watermark") + FUNCTION = "watermarkgen" + + + def watermarkgen(self, audio, add_watermark, key): + watermarker = self.load_watermarker(device=self.device) + audio_array, sample_rate = self.load_audio(audio) + # 确保 audio_array 在正确的设备上 + audio_array = audio_array.to(self.device) + + if add_watermark: + key = self._parse_key(key) + audio_array, sample_rate = self.watermark(watermarker, audio_array, sample_rate, key) + + watermark = self.verify(watermarker, audio_array, sample_rate) + + # 返回前将音频数据移回 CPU + return ({"waveform": audio_array.unsqueeze(0).unsqueeze(0).cpu(), "sample_rate": sample_rate}, watermark) + + @torch.inference_mode() + def watermark(self, + watermarker: silentcipher.server.Model, + audio_array: torch.Tensor, + sample_rate: int, + watermark_key: list[int], + ) -> tuple[torch.Tensor, int]: + # 确保音频是单声道 + if len(audio_array.shape) > 1 and audio_array.shape[0] > 1: + audio_array = audio_array.mean(dim=0) + + audio_array = audio_array.to(self.device) + + audio_array_44khz = torchaudio.functional.resample( + audio_array, + orig_freq=sample_rate, + new_freq=44100 + ).to(self.device) + + # 确保音频形状正确 (应为一维张量) + if len(audio_array_44khz.shape) != 1: + audio_array_44khz = audio_array_44khz.reshape(-1) + + + try: + # 增加水印强度，降低message_sdr值使水印更明显 + encoded, _ = watermarker.encode_wav(audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=30) + + verify_result = watermarker.decode_wav(encoded, 44100, phase_shift_decoding=True) + + if not verify_result["status"]: + encoded, _ = watermarker.encode_wav(audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=25) + verify_result = watermarker.decode_wav(encoded, 44100, phase_shift_decoding=True) + except Exception as e: + return audio_array, sample_rate + + # 如果需要，重采样回原始采样率 + output_sample_rate = min(44100, sample_rate) + if output_sample_rate != 44100: + encoded = torchaudio.functional.resample( + encoded, + orig_freq=44100, + new_freq=output_sample_rate + ).to(self.device) + + return encoded, output_sample_rate + + @torch.inference_mode() + def verify(self, + watermarker: silentcipher.server.Model, + watermarked_audio: torch.Tensor, + sample_rate: int, + ) -> str: + if len(watermarked_audio.shape) > 1 and watermarked_audio.shape[0] > 1: + watermarked_audio = watermarked_audio.mean(dim=0) + + if sample_rate != 44100: + watermarked_audio_44khz = torchaudio.functional.resample( + watermarked_audio, + orig_freq=sample_rate, + new_freq=44100 + ).to(self.device) + else: + watermarked_audio_44khz = watermarked_audio.to(self.device) + + if len(watermarked_audio_44khz.shape) != 1: + watermarked_audio_44khz = watermarked_audio_44khz.reshape(-1) + + + # 尝试不同的解码参数 + # 1. 使用相位偏移解码 + result_phase = watermarker.decode_wav(watermarked_audio_44khz, 44100, phase_shift_decoding=True) + + # 2. 不使用相位偏移解码 + result_no_phase = watermarker.decode_wav(watermarked_audio_44khz, 44100, phase_shift_decoding=False) + + # 使用两种方法中任一种成功的结果 + if result_phase["status"]: + watermark = "Watermarked:" + str(result_phase["messages"][0]) + elif result_no_phase["status"]: + watermark = "Watermarked:" + str(result_no_phase["messages"][0]) + else: + watermark = "No watermarked" + + return watermark + + + def load_watermarker(self, device: str = "cuda") -> silentcipher.server.Model: + ckpt_path = os.path.join(models_dir, "TTS", "SilentCipher", "44_1_khz", "73999_iteration") + config_path = os.path.join(models_dir, ckpt_path, "hparams.yaml") + model = silentcipher.get_model( + model_type="44.1k", + ckpt_path=ckpt_path, + config_path=config_path, + device=device, + ) + return model + + + def _parse_key(self, key_string): + """Helper function to safely parse the key.""" + try: + key = ast.literal_eval(key_string) + return key + except (ValueError, SyntaxError) as e: + raise + + + def load_audio(self, audio) -> tuple[torch.Tensor, int]: + waveform = audio["waveform"].squeeze(0) + audio_array = waveform.mean(dim=0) + + sample_rate = audio["sample_rate"] + + return audio_array, int(sample_rate) + + +@dataclass +class Segment: + speaker: int + text: str + # (num_samples,), sample_rate = 24_000 + audio: torch.Tensor + +SEGMENTS = [] +SPEAKERS = [] + +class Generator: + # 添加类变量用于缓存 + _cached_llama3_tokenizer = None + _cached_mimi = None + + def __init__( + self, + model: Model, + device: str = "cuda", + ): + self._model = model + self._model.setup_caches(1) + self.device = device + + self._text_tokenizer = self.load_llama3_tokenizer() + + mimi = self.load_mimi() + self._audio_tokenizer = mimi + self.sample_rate = mimi.sample_rate + + def load_llama3_tokenizer(self): + """ + https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992 + """ + if Generator._cached_llama3_tokenizer is not None: + return Generator._cached_llama3_tokenizer + + llama_path = os.path.join(models_dir, "LLM", "Llama-3.2-1B") + tokenizer = AutoTokenizer.from_pretrained(llama_path) + bos = tokenizer.bos_token + eos = tokenizer.eos_token + tokenizer._tokenizer.post_processor = TemplateProcessing( + single=f"{bos}:0 $A:0 {eos}:0", + pair=f"{bos}:0 $A:0 {eos}:0 {bos}:1 $B:1 {eos}:1", + special_tokens=[(f"{bos}", tokenizer.bos_token_id), (f"{eos}", tokenizer.eos_token_id)], + ) + Generator._cached_llama3_tokenizer = tokenizer + return tokenizer + + def load_mimi(self): + if Generator._cached_mimi is not None: + return Generator._cached_mimi + + mimi_path = os.path.join(models_dir, "TTS", "moshiko-pytorch-bf16", loaders.MIMI_NAME) + mimi = loaders.get_mimi(mimi_path, device=self.device) + mimi.set_num_codebooks(32) + Generator._cached_mimi = mimi + return mimi + + def _tokenize_text_segment(self, text: str, speaker: int) -> Tuple[torch.Tensor, torch.Tensor]: + frame_tokens = [] + frame_masks = [] + + text_tokens = self._text_tokenizer.encode(f"[{speaker}]{text}") + text_frame = torch.zeros(len(text_tokens), 33).long() + text_frame_mask = torch.zeros(len(text_tokens), 33).bool() + text_frame[:, -1] = torch.tensor(text_tokens) + text_frame_mask[:, -1] = True + + frame_tokens.append(text_frame.to(self.device)) + frame_masks.append(text_frame_mask.to(self.device)) + + return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0) + + def _tokenize_audio(self, audio: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + frame_tokens = [] + frame_masks = [] + + # (K, T) + audio = audio.to(self.device) + audio_tokens = self._audio_tokenizer.encode(audio.unsqueeze(0).unsqueeze(0))[0] + # add EOS frame + eos_frame = torch.zeros(audio_tokens.size(0), 1).to(self.device) + audio_tokens = torch.cat([audio_tokens, eos_frame], dim=1) + + audio_frame = torch.zeros(audio_tokens.size(1), 33).long().to(self.device) + audio_frame_mask = torch.zeros(audio_tokens.size(1), 33).bool().to(self.device) + audio_frame[:, :-1] = audio_tokens.transpose(0, 1) + audio_frame_mask[:, :-1] = True + + frame_tokens.append(audio_frame) + frame_masks.append(audio_frame_mask) + + return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0) + + def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Returns: + (seq_len, 33), (seq_len, 33) + """ + text_tokens, text_masks = self._tokenize_text_segment(segment.text, segment.speaker) + audio_tokens, audio_masks = self._tokenize_audio(segment.audio) + + return torch.cat([text_tokens, audio_tokens], dim=0), torch.cat([text_masks, audio_masks], dim=0) + + @torch.inference_mode() + def generate( + self, + text: str, + speaker: int, + context: List[Segment], + max_audio_length_ms: float = 90_000, + temperature: float = 0.9, + topk: int = 50, + ) -> torch.Tensor: + self._model.reset_caches() + + max_audio_frames = int(max_audio_length_ms / 80) + tokens, tokens_mask = [], [] + for segment in context: + segment_tokens, segment_tokens_mask = self._tokenize_segment(segment) + tokens.append(segment_tokens) + tokens_mask.append(segment_tokens_mask) + + gen_segment_tokens, gen_segment_tokens_mask = self._tokenize_text_segment(text, speaker) + tokens.append(gen_segment_tokens) + tokens_mask.append(gen_segment_tokens_mask) + + prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device) + prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device) + + samples = [] + curr_tokens = prompt_tokens.unsqueeze(0) + curr_tokens_mask = prompt_tokens_mask.unsqueeze(0) + curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device) + + max_seq_len = 2048 - max_audio_frames + if curr_tokens.size(1) >= max_seq_len: + raise ValueError(f"Inputs too long, must be below max_seq_len - max_audio_frames: {max_seq_len}") + + for _ in range(max_audio_frames): + sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk) + if torch.all(sample == 0): + break # eos + + samples.append(sample) + + curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1) + curr_tokens_mask = torch.cat( + [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1 + ).unsqueeze(1) + curr_pos = curr_pos[:, -1:] + 1 + + audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0) + + return audio + + +class MultiLinePromptCSM: + @classmethod + def INPUT_TYPES(cls): + + return { + "required": { + "multi_line_prompt": ("STRING", { + "multiline": True, + "default": ""}), + }, + } + + CATEGORY = "MW_CSM" + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("prompt",) + FUNCTION = "promptgen" + + def promptgen(self, multi_line_prompt: str): + return (multi_line_prompt.strip(),) + + +class CSMDialogRun: + # 添加类变量用于缓存 + _cached_csm_1b = None + _cached_generator = None + + if torch.backends.mps.is_available(): + device = "mps" + elif torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + @classmethod + def INPUT_TYPES(s): + return {"required": { + "text": ("STRING",), + "unload_speakers": ("BOOLEAN",{ "default": False}), + }, + "optional": { + "prompt0": ("STRING",), + "prompt1": ("STRING",), + "prompt2": ("STRING",), + "prompt3": ("STRING",), + "audio0": ("AUDIO",), + "audio1": ("AUDIO",), + "audio2": ("AUDIO",), + "audio3": ("AUDIO",), + "who_will_speak": ("INT", {"default": 0, "min": 0, "max": 9, "step": 1,}), + } + } + + + CATEGORY = "MW_CSM" + RETURN_TYPES = ("AUDIO", "STRING") + RETURN_NAMES = ("audio", "prompt") + FUNCTION = "run" + + + def run(self, text, unload_speakers, prompt0="", prompt1="", prompt2="", prompt3="", audio0=None, audio1=None, audio2=None, audio3=None, who_will_speak=1): + generator = self.load_csm_1b() + global SEGMENTS, SPEAKERS + if unload_speakers: + SEGMENTS.clear() + SPEAKERS.clear() + segments = [] + + for i in range(4): + prompt = locals()[f"prompt{i}"] + audio = locals()[f"audio{i}"] + # print(f"prompt{i}: {prompt}→audio{i}") + if audio is not None: + audio_tensor = audio["waveform"].squeeze(0).mean(dim=0) + sample_rate = int(audio["sample_rate"]) + # print(f"audio{i} sample_rate: {sample_rate}") + audio_tensor = torchaudio.functional.resample( + audio_tensor.squeeze(0), + orig_freq=sample_rate, + new_freq=generator.sample_rate + ) + else: + audio_tensor = None + + speaker, prompt = self.get_speaker_text(prompt) + + segment = self.get_segment(speaker, prompt, audio_tensor) + if segment is not None: + SEGMENTS.append(segment) + SPEAKERS.append(speaker) + + if SEGMENTS: + SEGMENTS.extend(segments) + if len(SEGMENTS) > 6: + SEGMENTS = SEGMENTS[-6:] + SPEAKERS = SPEAKERS[-6:] + if who_will_speak not in SPEAKERS: + raise ValueError(f"The speaker {who_will_speak} not found or cleared, up to 6 recent speakers saved.") + + # print(f"使用 {len(segments)} 个上下文段落生成音频，说话者: {will_speaker}") + audio = generator.generate( + text=text, + speaker=who_will_speak, + context=SEGMENTS, + max_audio_length_ms=10_000, + ) + out_prompt = str(who_will_speak) + ": " + text + else: + # print(f"无上下文生成音频，说话者: 0") + audio = generator.generate( + text=text, + speaker=0, + context=[], + max_audio_length_ms=10_000, + ) + out_prompt = "0: " + text + return ({"waveform": audio.unsqueeze(0).unsqueeze(0).cpu(), "sample_rate": generator.sample_rate}, out_prompt) + + def get_speaker_text(self, text): + import re + if text.strip() != "": + st = [i.strip() for i in re.split('[:：]', text, 1)] + if len(st) == 2: + speaker, prompt = st + return int(speaker), prompt + else: + raise ValueError("Invalid text format") + else: + return None, None + + def get_segment(self, speaker, text, audio): + if speaker is not None: + if audio is not None: + return Segment(speaker=speaker, text=text, audio=audio) + else: + raise ValueError(f"{text}: Audio is required") + else: + return None + + def load_csm_1b(self) -> Generator: + if CSMDialogRun._cached_generator is not None: + return CSMDialogRun._cached_generator + + if CSMDialogRun._cached_csm_1b is None: + csm_1b_path = os.path.join(models_dir, "TTS", "csm-1b") + config_path = os.path.join(csm_1b_path, "config.json") + import json + with open(config_path, 'r') as f: + config = json.load(f)["args"] + + configs = ModelArgs(backbone_flavor = config["backbone_flavor"], + decoder_flavor = config["decoder_flavor"], + text_vocab_size = config["text_vocab_size"], + audio_vocab_size = config["audio_vocab_size"], + audio_num_codebooks = config["audio_num_codebooks"]) + model = Model.from_pretrained(csm_1b_path, config=configs) + model.to(device=self.device, dtype=torch.bfloat16) + CSMDialogRun._cached_csm_1b = model + + generator = Generator(CSMDialogRun._cached_csm_1b, device=self.device) + CSMDialogRun._cached_generator = generator + return generator + +from .MWAudioRecorderCSM import AudioRecorderCSM + +NODE_CLASS_MAPPINGS = { + "AddWatermark": AddWatermark, + "CSMDialogRun": CSMDialogRun, + "MultiLinePromptCSM": MultiLinePromptCSM, + "AudioRecorderCSM": AudioRecorderCSM, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AddWatermark": "Add Watermark", + "CSMDialogRun": "CSM Dialog Run", + "MultiLinePromptCSM": "Multi Line Prompt", + "AudioRecorderCSM": "MW Audio Recorder", +} \ No newline at end of file diff --git a/MWAudioRecorderCSM.py b/MWAudioRecorderCSM.py new file mode 100644 index 0000000..64c68d3 --- /dev/null +++ b/MWAudioRecorderCSM.py @@ -0,0 +1,138 @@ +import numpy as np +import torch +import time +import librosa +import sounddevice as sd +from scipy import ndimage +from comfy.utils import ProgressBar + + + +class AudioRecorderCSM: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # 触发控制 + "trigger": ("BOOLEAN", {"default": False}), + # 录音时长 + "record_sec": ("INT", { + "default": 5, + "min": 1, + "max": 60, + "step": 1 # 整数秒递增 + }), + "sample_rate": (["16000", "44100", "48000"], { # 限定标准采样率 + "default": "48000" + }), + "n_fft": ("INT", { # 限定为2的幂次方 + "default": 2048, + "min": 512, + "max": 4096, + "step": 512 # 只能选择512,1024,1536,2048,...4096 + }), + "sensitivity": ("FLOAT", { # 灵敏度精确控制 + "default": 1.2, + "min": 0.5, + "max": 3.0, + "step": 0.1 # 0.1步进 + }), + "smooth": ("INT", { # 确保为奇数 + "default": 5, + "min": 1, + "max": 11, + "step": 2 # 生成1,3,5,7,9,11 + }), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + }, + "optional": { + "interlocutor": ("AUDIO",), + }, + } + + RETURN_TYPES = ("AUDIO", "AUDIO") + RETURN_NAMES = ("audio", "interlocutor") + FUNCTION = "record_and_clean" + CATEGORY = "MW_CSM" + + def _stft(self, y, n_fft): + hop = n_fft // 4 + return librosa.stft(y, n_fft=n_fft, hop_length=hop, win_length=n_fft) + + def _istft(self, spec, n_fft): + hop = n_fft // 4 + return librosa.istft(spec, hop_length=hop, win_length=n_fft) + + def _calc_noise_profile(self, noise_clip, n_fft): + noise_spec = self._stft(noise_clip, n_fft) + return { + 'mean': np.mean(np.abs(noise_spec), axis=1, keepdims=True), + 'std': np.std(np.abs(noise_spec), axis=1, keepdims=True) + } + + def _spectral_gate(self, spec, noise_profile, sensitivity): + threshold = noise_profile['mean'] + sensitivity * noise_profile['std'] + return np.where(np.abs(spec) > threshold, spec, 0) + + def _smooth_mask(self, mask, kernel_size): + smoothed = ndimage.uniform_filter(mask, size=(kernel_size, kernel_size)) + return np.clip(smoothed * 1.2, 0, 1) # 增强边缘保留 + + def record_and_clean(self, trigger, record_sec, n_fft, sensitivity, smooth, sample_rate, interlocutor=None, seed=0): + if not trigger: + if interlocutor is not None: + return (None, interlocutor) + raise ValueError("No trigger received: Recording not opened.") + + sr = int(sample_rate) + final_audio = None + + try: + noise_clip = None + # 主录音 + # print(f"开始主录音 {record_sec}秒...") + main_rec = sd.rec(int(record_sec * sr), samplerate=sr, channels=1, dtype='float32') + pb = ProgressBar(record_sec) + for _ in range(record_sec * 2): + time.sleep(0.5) + pb.update(0.5) + sd.wait() + audio = main_rec.flatten() + + # 自动噪声检测 + if noise_clip is None: + # print("自动检测静默段作为噪声参考...") + energy = librosa.feature.rms(y=audio, frame_length=n_fft, hop_length=n_fft//4) + min_idx = np.argmin(energy) + start = min_idx * (n_fft//4) + noise_clip = audio[start:start + n_fft*2] + + # 降噪处理 + # print("进行频谱降噪...") + noise_profile = self._calc_noise_profile(noise_clip, n_fft) + spec = self._stft(audio, n_fft) + + # 多步骤处理 + mask = np.ones_like(spec) # 初始掩膜 + for _ in range(2): # 双重处理循环 + cleaned_spec = self._spectral_gate(spec, noise_profile, sensitivity) + mask = np.where(np.abs(cleaned_spec) > 0, 1, 0) + mask = self._smooth_mask(mask, smooth//2+1) + spec = spec * mask + + # 相位恢复重建 + processed = self._istft(spec * mask, n_fft) + + # 动态增益归一化 + peak = np.max(np.abs(processed)) + processed = processed * (0.99 / peak) if peak > 0 else processed + + # 格式转换 + waveform = torch.from_numpy(processed).float().unsqueeze(0).unsqueeze(0) + final_audio = {"waveform": waveform, "sample_rate": sr} + + except Exception as e: + print(f"Recording/processing failed: {str(e)}") + raise + + return (final_audio, interlocutor) \ No newline at end of file diff --git a/README-CN.md b/README-CN.md new file mode 100644 index 0000000..f64e461 --- /dev/null +++ b/README-CN.md @@ -0,0 +1,49 @@ +[中文](README-CN.md)|[English](README.md) + +# ComfyUI Node for CSM + +CSM（Conversational Speech Model 会话语音模型）, 多人会话, 克隆声音, 然后可根据会话中语音的情绪变化等, 生成相应情绪变化的语音的模型. 遗憾的是目前只有英文可用. 该节点, 暂时支持同时 10 人会话. + +![](https://github.com/billwuhao/ComfyUI_CSM/blob/master/images/2025-03-18_15-38-45.png) + +可以将录音节点穿插在其中, 进行多人会话. + +![](https://github.com/billwuhao/ComfyUI_CSM/blob/master/images/2025-03-18_19-01-15.png) + +还支持音频水印检测(自动检测水印) 和 音频添加加密水印. + +![](https://github.com/billwuhao/ComfyUI_CSM/blob/master/images/2025-03-18_14-43-49.png) + +## 📣 更新 + +[2025-03-18]⚒️: 发布版本 v1.0.0. + +节点使用方法详解示例工作流: [example_workflows](https://github.com/billwuhao/ComfyUI_CSM/blob/master/example_workflows) + +## 安装 + +``` +cd ComfyUI/custom_nodes +git clone https://github.com/billwuhao/ComfyUI_CSM.git +cd ComfyUI_CSM +pip install -r requirements.txt + +# python_embeded +./python_embeded/python.exe -m pip install -r requirements.txt +``` + +## 模型下载 + +- [csm-1b](https://huggingface.co/sesame/csm-1b/tree/main): `config.json` 和 `model.safetensors` 下载放到 `ComfyUI/models/TTS/csm-1b` 目录下. + +- [moshiko-pytorch-bf16](https://huggingface.co/kyutai/moshiko-pytorch-bf16/tree/main): `tokenizer-e351c8d8-checkpoint125.safetensors` 下载放到 `ComfyUI/models/TTS/moshiko-pytorch-bf16` 目录下. + +- [SilentCipher](https://huggingface.co/Sony/SilentCipher/tree/main/44_1_khz/73999_iteration): 全部模型下载放到 `ComfyUI\models\TTS\SilentCipher\44_1_khz\73999_iteration` 目录下. + +- [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B/tree/main): 除了 `original` 目录, 其他全部下载放到 `ComfyUI\models\LLM\Llama-3.2-1B` 目录下. + +## 鸣谢 + +[csm](https://github.com/SesameAILabs/csm) + +感谢 SesameAILabs 团队的卓越的工作👍. \ No newline at end of file diff --git a/README.md b/README.md index bced0ca..d2bacc3 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,49 @@ -# ComfyUI_CSM -ComfyUI node of Conversational Speech Model (CSM). +[中文](README-CN.md)|[English](README.md) + +# ComfyUI Node for CSM + +CSM (Conversational Speech Model), a model that supports multi-person conversations, voice cloning, and generates speech with corresponding emotional changes based on the emotional changes in the conversation. Unfortunately, it is currently only available in English. This node temporarily supports simultaneous conversations with up to 10 people. + +![](https://github.com/billwuhao/ComfyUI_CSM/blob/master/images/2025-03-18_15-38-45.png) + +Recording nodes can be interspersed within to create multi-person conversations. + +![](https://github.com/billwuhao/ComfyUI_CSM/blob/master/images/2025-03-18_19-01-15.png) + +It also supports audio watermark detection (automatic watermark detection) and audio adding encrypted watermarks. + +![](https://github.com/billwuhao/ComfyUI_CSM/blob/master/images/2025-03-18_14-43-49.png) + +## 📣 Updates + +[2025-03-18]⚒️: Released version v1.0.0. + +Detailed node usage example workflows: [example_workflows](https://github.com/billwuhao/ComfyUI_CSM/blob/master/example_workflows) + +## Installation + +``` +cd ComfyUI/custom_nodes +git clone https://github.com/billwuhao/ComfyUI_CSM.git +cd ComfyUI_CSM +pip install -r requirements.txt + +# python_embeded +./python_embeded/python.exe -m pip install -r requirements.txt +``` + +## Model Download + +- [csm-1b](https://huggingface.co/sesame/csm-1b/tree/main): Download `config.json` and `model.safetensors` and place them in the `ComfyUI/models/TTS/csm-1b` directory. + +- [moshiko-pytorch-bf16](https://huggingface.co/kyutai/moshiko-pytorch-bf16/tree/main): Download `tokenizer-e351c8d8-checkpoint125.safetensors` and place it in the `ComfyUI/models/TTS/moshiko-pytorch-bf16` directory. + +- [SilentCipher](https://huggingface.co/Sony/SilentCipher/tree/main/44_1_khz/73999_iteration): Download all models and place them in the `ComfyUI\models\TTS\SilentCipher\44_1_khz\73999_iteration` directory. + +- [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B/tree/main): Download everything except the `original` directory and place it in the `ComfyUI\models\LLM\Llama-3.2-1B` directory. + +## Acknowledgements + +[csm](https://github.com/SesameAILabs/csm) + +Thanks to the SesameAILabs team for their excellent work 👍. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..bd91062 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .CSMNode import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/example_workflows/Multi-User Dialogue.json b/example_workflows/Multi-User Dialogue.json new file mode 100644 index 0000000..0eac9d4 --- /dev/null +++ b/example_workflows/Multi-User Dialogue.json @@ -0,0 +1 @@ +{"last_node_id":37,"last_link_id":87,"nodes":[{"id":2,"type":"PreviewAudio","pos":[689.6212158203125,-131.0446319580078],"size":[315,88],"flags":{},"order":15,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":70}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":6,"type":"PreviewAudio","pos":[1794.6282958984375,-127.8306884765625],"size":[315,88],"flags":{},"order":20,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":76}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":3,"type":"MultiLinePromptCSM","pos":[669.1260375976562,423.7393798828125],"size":[336.4093322753906,175.87789916992188],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[65],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["hello, what are you doing "]},{"id":11,"type":"MultiLinePromptCSM","pos":[11.868524551391602,48.75131607055664],"size":[237.15884399414062,166.07469177246094],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[66],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["0: The sky above the port was the color of television, tuned to a dead channel"]},{"id":12,"type":"MultiLinePromptCSM","pos":[352.0327453613281,41.29263687133789],"size":[237.15884399414062,166.07469177246094],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[67],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["1: And Lake turned round upon me, a little abruptly, his odd yellowish eyes, a little like those of the sea eagle, and the ghost of his smile that flickered on his singularly pale face, with a stern and insidious look, confronted me."]},{"id":15,"type":"MultiLinePromptCSM","pos":[16.676734924316406,535.7870483398438],"size":[237.15884399414062,166.07469177246094],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[68],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["4: like a super Mario level. Like it's very like high detail. And like, once you get into the park, it just like, everything looks like a computer game and they have all these, like, you know, if, if there's like a, you know, like in a Mario game, they will have like a question block. And if you like, you know, punch it, a coin will come out. So like everyone, when they come into the park, they get like this little bracelet and then you can go punching question blocks around.."]},{"id":20,"type":"MultiLinePromptCSM","pos":[354.8899230957031,537.6543579101562],"size":[237.15884399414062,166.07469177246094],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[69],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["5: Suddenly I was back in the old days Before you felt we ought to drift apart. It was some trick-the way your eyebrows raise. "]},{"id":23,"type":"easy showAnything","pos":[1075.70263671875,364.15948486328125],"size":[210,165.15225219726562],"flags":{},"order":16,"mode":0,"inputs":[{"name":"anything","localized_name":"anything","type":"*","shape":7,"link":71}],"outputs":[{"name":"output","localized_name":"output","type":"*","links":[79],"slot_index":0}],"properties":{"cnr_id":"comfyui-easy-use","ver":"3e84b8cd77719341adc71b08b5789bdb07b1a543","Node name for S&R":"easy showAnything"},"widgets_values":["4: hello, what are you doing"]},{"id":8,"type":"MultiLinePromptCSM","pos":[1369.939697265625,421.773681640625],"size":[342.322021484375,170.18701171875],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[72],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["Nice to meet you all "]},{"id":4,"type":"PreviewAudio","pos":[1384.184814453125,-130.83311462402344],"size":[315,88],"flags":{},"order":18,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":73}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":33,"type":"CSMDialogRun","pos":[1381.4351806640625,6.7029709815979],"size":[306.3296203613281,362],"flags":{},"order":17,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":77},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":80},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":72},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":78},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":79},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":null},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[73,82],"slot_index":0},{"name":"prompt","localized_name":"prompt","type":"STRING","links":[83]}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",0]},{"id":9,"type":"MultiLinePromptCSM","pos":[1798.8135986328125,425.1792297363281],"size":[305.42041015625,173.27102661132812],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[75],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["What a nice day today "]},{"id":35,"type":"PreviewAudio","pos":[2181.03515625,-126.64925384521484],"size":[315,88],"flags":{},"order":22,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":84}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":10,"type":"LoadAudio","pos":[14.130395889282227,-144.62242126464844],"size":[238.1028289794922,136],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"AUDIO","localized_name":"AUDIO","type":"AUDIO","links":[61]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"LoadAudio"},"widgets_values":["af_heart_0.wav","",""]},{"id":25,"type":"MultiLinePromptCSM","pos":[1066.255859375,69.8873519897461],"size":[237.15884399414062,166.07469177246094],"flags":{},"order":8,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[78],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["6: All passed so quickly, there was so much going on around him, the Tree quite forgot to look to himself."]},{"id":37,"type":"MultiLinePromptCSM","pos":[2185.220458984375,426.36065673828125],"size":[305.42041015625,173.27102661132812],"flags":{},"order":9,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[85],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["Goodbye"]},{"id":14,"type":"LoadAudio","pos":[12.917945861816406,342.4131774902344],"size":[238.1028289794922,136],"flags":{},"order":10,"mode":0,"inputs":[],"outputs":[{"name":"AUDIO","localized_name":"AUDIO","type":"AUDIO","links":[63],"slot_index":0}],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"LoadAudio"},"widgets_values":["conversational_b.wav","",""]},{"id":32,"type":"CSMDialogRun","pos":[680.8510131835938,21.543092727661133],"size":[315,362],"flags":{},"order":14,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":61},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":62},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":63},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":64},{"name":"text","type":"STRING","widget":{"name":"text"},"link":65},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":66},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":67},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":68},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":69}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[70,80],"slot_index":0},{"name":"prompt","localized_name":"prompt","type":"STRING","links":[71]}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",4]},{"id":24,"type":"LoadAudio","pos":[1068.5177001953125,-123.4864273071289],"size":[238.1028289794922,136],"flags":{},"order":11,"mode":0,"inputs":[],"outputs":[{"name":"AUDIO","localized_name":"AUDIO","type":"AUDIO","links":[77],"slot_index":0}],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"LoadAudio"},"widgets_values":["read_speech_c.wav","",""]},{"id":34,"type":"CSMDialogRun","pos":[1797.8992919921875,29.272401809692383],"size":[315,362],"flags":{},"order":19,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":82},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":null},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":75},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":83},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":null},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":null},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[76,87],"slot_index":0},{"name":"prompt","localized_name":"prompt","type":"STRING","links":[86],"slot_index":1}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",5]},{"id":13,"type":"LoadAudio","pos":[351.2843933105469,-147.56602478027344],"size":[238.1028289794922,136],"flags":{},"order":12,"mode":0,"inputs":[],"outputs":[{"name":"AUDIO","localized_name":"AUDIO","type":"AUDIO","links":[62],"slot_index":0}],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"LoadAudio"},"widgets_values":["read_speech_a.wav","",""]},{"id":19,"type":"LoadAudio","pos":[355.6466369628906,344.2802429199219],"size":[238.1028289794922,136],"flags":{},"order":13,"mode":0,"inputs":[],"outputs":[{"name":"AUDIO","localized_name":"AUDIO","type":"AUDIO","links":[64],"slot_index":0}],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"LoadAudio"},"widgets_values":["read_speech_d.wav","",""]},{"id":36,"type":"CSMDialogRun","pos":[2184.30615234375,30.453838348388672],"size":[315,362],"flags":{},"order":21,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":87},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":null},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":85},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":86},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":null},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":null},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[84]},{"name":"prompt","localized_name":"prompt","type":"STRING","links":null}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",6]}],"links":[[61,10,0,32,0,"AUDIO"],[62,13,0,32,1,"AUDIO"],[63,14,0,32,2,"AUDIO"],[64,19,0,32,3,"AUDIO"],[65,3,0,32,4,"STRING"],[66,11,0,32,5,"STRING"],[67,12,0,32,6,"STRING"],[68,15,0,32,7,"STRING"],[69,20,0,32,8,"STRING"],[70,32,0,2,0,"AUDIO"],[71,32,1,23,0,"*"],[72,8,0,33,4,"STRING"],[73,33,0,4,0,"AUDIO"],[75,9,0,34,4,"STRING"],[76,34,0,6,0,"AUDIO"],[77,24,0,33,0,"AUDIO"],[78,25,0,33,5,"STRING"],[79,23,0,33,6,"STRING"],[80,32,0,33,1,"AUDIO"],[82,33,0,34,0,"AUDIO"],[83,33,1,34,5,"STRING"],[84,36,0,35,0,"AUDIO"],[85,37,0,36,4,"STRING"],[86,34,1,36,5,"STRING"],[87,34,0,36,0,"AUDIO"]],"groups":[{"id":1,"title":"Group1","bounding":[659.1260375976562,-209.1600799560547,354.9040832519531,821.78759765625],"color":"#3f789e","font_size":24,"flags":{}},{"id":2,"title":"Group2","bounding":[1359.939697265625,-205.93833923339844,378.8785400390625,807.8990478515625],"color":"#3f789e","font_size":24,"flags":{}},{"id":3,"title":"Group3","bounding":[1784.6280517578125,-201.43064880371094,338.27099609375,806.8706665039062],"color":"#3f789e","font_size":24,"flags":{}},{"id":4,"title":"speaker0","bounding":[1.8685262203216553,-218.22247314453125,260.3647155761719,443.0484619140625],"color":"#3f789e","font_size":24,"flags":{}},{"id":5,"title":"speaker1","bounding":[339.0224914550781,-221.166015625,260.3647155761719,443.0484619140625],"color":"#3f789e","font_size":24,"flags":{}},{"id":6,"title":"speaker4","bounding":[6.6767354011535645,268.81317138671875,260.3647155761719,443.0484619140625],"color":"#3f789e","font_size":24,"flags":{}},{"id":7,"title":"speaker5","bounding":[343.38482666015625,270.68023681640625,260.3647155761719,443.0484619140625],"color":"#3f789e","font_size":24,"flags":{}},{"id":8,"title":"speaker6","bounding":[1056.255859375,-197.08642578125,260.36474609375,443.0484924316406],"color":"#3f789e","font_size":24,"flags":{}},{"id":9,"title":"Group4","bounding":[2171.034912109375,-200.24923706054688,338.27099609375,806.8706665039062],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"ds":{"scale":0.620921323059155,"offset":[124.20233837409961,298.01713535129545]},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/example_workflows/Recording dialogue.json b/example_workflows/Recording dialogue.json new file mode 100644 index 0000000..d672eed --- /dev/null +++ b/example_workflows/Recording dialogue.json @@ -0,0 +1 @@ +{"last_node_id":42,"last_link_id":95,"nodes":[{"id":6,"type":"AudioRecorderCSM","pos":[809.9296264648438,340.3097839355469],"size":[224.876708984375,259.0841064453125],"flags":{},"order":0,"mode":0,"inputs":[{"name":"interlocutor","localized_name":"interlocutor","type":"AUDIO","shape":7,"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[72],"slot_index":0},{"name":"interlocutor","localized_name":"interlocutor","type":"AUDIO","links":null}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"AudioRecorderCSM"},"widgets_values":[true,7,"48000",2560,0.5,1,1111813690525763,"fixed"]},{"id":8,"type":"PreviewAudio","pos":[1158.3402099609375,-74.12663269042969],"size":[315,88],"flags":{},"order":10,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":76}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":10,"type":"MultiLinePromptCSM","pos":[812.62646484375,77.07012939453125],"size":[219.40182495117188,129.62608337402344],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[74],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["0: like revising for an exam I'd have to try and like keep up the momentum because I'd start really early I'd be like okay I'm gonna start revising now and then like you're revising for ages and then I just like start losing steam I didn't do that for the exam we had recently to be fair that was a more of a last minute scenario but like yeah I'm trying to like yeah I noticed this yesterday that like Mondays I sort of start the day with this not like a panic but like a"]},{"id":11,"type":"MultiLinePromptCSM","pos":[813.6771850585938,649.297119140625],"size":[226.24343872070312,178.88548278808594],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[75],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["1: The sky above the port was the color of television, tuned to a dead channel"]},{"id":12,"type":"MultiLinePromptCSM","pos":[1146.166015625,519.1640625],"size":[389.1127014160156,281.9698181152344],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[73],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["Hello, what are you doing"]},{"id":15,"type":"easy showAnything","pos":[1366.4840087890625,155.53948974609375],"size":[210,229.8733367919922],"flags":{},"order":12,"mode":0,"inputs":[{"name":"anything","localized_name":"anything","type":"*","shape":7,"link":87}],"outputs":[{"name":"output","localized_name":"output","type":"*","links":[],"slot_index":0}],"properties":{"cnr_id":"comfyui-easy-use","ver":"3e84b8cd77719341adc71b08b5789bdb07b1a543","Node name for S&R":"easy showAnything"},"widgets_values":["0: Hello, what are you doing"]},{"id":20,"type":"easy showAnything","pos":[2179.247802734375,88.58969116210938],"size":[216.02073669433594,219.33731079101562],"flags":{},"order":15,"mode":0,"inputs":[{"name":"anything","localized_name":"anything","type":"*","shape":7,"link":89}],"outputs":[{"name":"output","localized_name":"output","type":"*","links":[93],"slot_index":0}],"properties":{"cnr_id":"comfyui-easy-use","ver":"3e84b8cd77719341adc71b08b5789bdb07b1a543","Node name for S&R":"easy showAnything"},"widgets_values":["1: Nice to meet you"]},{"id":21,"type":"MultiLinePromptCSM","pos":[1960.1605224609375,501.2001953125],"size":[407.1747131347656,260.8977355957031],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[79],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["Nice to meet you"]},{"id":27,"type":"MultiLinePromptCSM","pos":[2467.9541015625,514.046630859375],"size":[435.77239990234375,239.8255615234375],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[82],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["Have you seen a cute black kitten "]},{"id":28,"type":"easy showAnything","pos":[2697.577392578125,56.28163146972656],"size":[210,229.8733367919922],"flags":{},"order":18,"mode":0,"inputs":[{"name":"anything","localized_name":"anything","type":"*","shape":7,"link":90}],"outputs":[{"name":"output","localized_name":"output","type":"*","links":[94],"slot_index":0}],"properties":{"cnr_id":"comfyui-easy-use","ver":"3e84b8cd77719341adc71b08b5789bdb07b1a543","Node name for S&R":"easy showAnything"},"widgets_values":["5: Have you seen a cute black kitten"]},{"id":33,"type":"MultiLinePromptCSM","pos":[2986.296875,505.9626159667969],"size":[434.2672424316406,235.31005859375],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[84],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["Goodbye, see you next time "]},{"id":34,"type":"easy showAnything","pos":[3203.87890625,79.80587768554688],"size":[210,229.8733367919922],"flags":{},"order":21,"mode":0,"inputs":[{"name":"anything","localized_name":"anything","type":"*","shape":7,"link":91}],"outputs":[{"name":"output","localized_name":"output","type":"*","links":[],"slot_index":0}],"properties":{"cnr_id":"comfyui-easy-use","ver":"3e84b8cd77719341adc71b08b5789bdb07b1a543","Node name for S&R":"easy showAnything"},"widgets_values":["0: Goodbye, see you next time"]},{"id":37,"type":"MultiLinePromptCSM","pos":[1655.0902099609375,221.94337463378906],"size":[223.89755249023438,317.401611328125],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[80],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["5: All passed so quickly, there was so much going on around him, the Tree quite forgot to look to himself."]},{"id":39,"type":"CSMDialogRun","pos":[1132.1058349609375,65.41092681884766],"size":[210,366.5155029296875],"flags":{},"order":9,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":71},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":72},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":73},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":74},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":75},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":null},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[76,77]},{"name":"prompt","localized_name":"prompt","type":"STRING","links":[87],"slot_index":1}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",0]},{"id":40,"type":"CSMDialogRun","pos":[1945.5535888671875,35.67900085449219],"size":[210,362],"flags":{},"order":13,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":78},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":null},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":79},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":80},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":null},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":null},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[81,92],"slot_index":0},{"name":"prompt","localized_name":"prompt","type":"STRING","links":[89],"slot_index":1}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",1]},{"id":41,"type":"CSMDialogRun","pos":[2460.873291015625,25.948413848876953],"size":[210,362],"flags":{},"order":16,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":92},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":null},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":82},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":93},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":null},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":null},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[83,95],"slot_index":0},{"name":"prompt","localized_name":"prompt","type":"STRING","links":[90],"slot_index":1}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",5]},{"id":42,"type":"CSMDialogRun","pos":[2967.174560546875,32.91583251953125],"size":[210,363.5050964355469],"flags":{},"order":19,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":95},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":null},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":84},{"name":"prompt0","type":"STRING","shape":7,"widget":{"name":"prompt0"},"link":94},{"name":"prompt1","type":"STRING","shape":7,"widget":{"name":"prompt1"},"link":null},{"name":"prompt2","type":"STRING","shape":7,"widget":{"name":"prompt2"},"link":null},{"name":"prompt3","type":"STRING","shape":7,"widget":{"name":"prompt3"},"link":null}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[85]},{"name":"prompt","localized_name":"prompt","type":"STRING","links":[91],"slot_index":1}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["",false,"","","","",0]},{"id":16,"type":"AudioRecorderCSM","pos":[1651.752685546875,-88.54882049560547],"size":[232.4023895263672,253.06346130371094],"flags":{},"order":11,"mode":0,"inputs":[{"name":"interlocutor","localized_name":"interlocutor","type":"AUDIO","shape":7,"link":77}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[78],"slot_index":0},{"name":"interlocutor","localized_name":"interlocutor","type":"AUDIO","links":[],"slot_index":1}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"AudioRecorderCSM"},"widgets_values":[true,10,"48000",2560,0.5,1,1095838717141322,"fixed"]},{"id":18,"type":"PreviewAudio","pos":[1968.7774658203125,-105.9108657836914],"size":[315,88],"flags":{},"order":14,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":81}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":26,"type":"PreviewAudio","pos":[2490.11767578125,-105.10552978515625],"size":[315,88],"flags":{},"order":17,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":83}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":32,"type":"PreviewAudio","pos":[3003.94482421875,-96.63292694091797],"size":[315,88],"flags":{},"order":20,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":85}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":9,"type":"LoadAudio","pos":[815.822998046875,-106.07005310058594],"size":[210,136],"flags":{},"order":8,"mode":0,"inputs":[],"outputs":[{"name":"AUDIO","localized_name":"AUDIO","type":"AUDIO","links":[71],"slot_index":0}],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"LoadAudio"},"widgets_values":["conversational_a.wav","",""]}],"links":[[71,9,0,39,0,"AUDIO"],[72,6,0,39,1,"AUDIO"],[73,12,0,39,4,"STRING"],[74,10,0,39,5,"STRING"],[75,11,0,39,6,"STRING"],[76,39,0,8,0,"AUDIO"],[77,39,0,16,0,"AUDIO"],[78,16,0,40,0,"AUDIO"],[79,21,0,40,4,"STRING"],[80,37,0,40,5,"STRING"],[81,40,0,18,0,"AUDIO"],[82,27,0,41,4,"STRING"],[83,41,0,26,0,"AUDIO"],[84,33,0,42,4,"STRING"],[85,42,0,32,0,"AUDIO"],[87,39,1,15,0,"*"],[89,40,1,20,0,"*"],[90,41,1,28,0,"*"],[91,42,1,34,0,"*"],[92,40,0,41,0,"AUDIO"],[93,20,0,41,5,"STRING"],[94,28,0,42,5,"STRING"],[95,41,0,42,0,"AUDIO"]],"groups":[{"id":1,"title":"speaker0","bounding":[796.4398193359375,-181.85504150390625,255.79464721679688,403.1885681152344],"color":"#3f789e","font_size":24,"flags":{}},{"id":2,"title":"speaker1","bounding":[799.9296264648438,266.70977783203125,258.5599365234375,587.8922119140625],"color":"#3f789e","font_size":24,"flags":{}},{"id":3,"title":"speaker5","bounding":[1635.732177734375,-168.1693115234375,262.6649169921875,754.8269653320312],"color":"#3f789e","font_size":24,"flags":{}},{"id":4,"title":"Group1","bounding":[1112.5125732421875,-165.78842163085938,475.4935302734375,987.1588134765625],"color":"#3f789e","font_size":24,"flags":{}},{"id":5,"title":"Group3","bounding":[1933.740478515625,-178.96359252929688,471.5273742675781,974.8699951171875],"color":"#3f789e","font_size":24,"flags":{}},{"id":6,"title":"Group4","bounding":[2446.049560546875,-173.64279174804688,476.04296875,949.2825317382812],"color":"#3f789e","font_size":24,"flags":{}},{"id":7,"title":"Group5","bounding":[2956.866455078125,-175.70620727539062,476.04278564453125,941.7568359375],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"ds":{"scale":0.620921323059155,"offset":[-732.4279397249601,280.5052479580308]},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/example_workflows/Text to speech and watermarking.json b/example_workflows/Text to speech and watermarking.json new file mode 100644 index 0000000..4432695 --- /dev/null +++ b/example_workflows/Text to speech and watermarking.json @@ -0,0 +1 @@ +{"last_node_id":15,"last_link_id":15,"nodes":[{"id":2,"type":"LoadAudio","pos":[-1.9409359693527222,-13.219098091125488],"size":[315,136],"flags":{},"order":0,"mode":4,"inputs":[],"outputs":[{"name":"AUDIO","localized_name":"AUDIO","type":"AUDIO","links":[9],"slot_index":0}],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"LoadAudio"},"widgets_values":["ComfyUI_00001_.flac","",""]},{"id":5,"type":"AddWatermark","pos":[371.0498352050781,-3.9667160511016846],"size":[315,102],"flags":{},"order":2,"mode":4,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":9}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[5],"slot_index":0},{"name":"watermark","localized_name":"watermark","type":"STRING","links":[6]}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"AddWatermark"},"widgets_values":[true,"[0, 211, 146, 56, 201]"]},{"id":4,"type":"PreviewAudio","pos":[749.3675537109375,-7.144290447235107],"size":[315,88],"flags":{},"order":4,"mode":4,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":5}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]},{"id":3,"type":"easy showAnything","pos":[762.9191284179688,143.41648864746094],"size":[316.5421142578125,88],"flags":{},"order":5,"mode":4,"inputs":[{"name":"anything","localized_name":"anything","type":"*","shape":7,"link":6}],"outputs":[{"name":"output","localized_name":"output","type":"*","links":null}],"properties":{"cnr_id":"comfyui-easy-use","ver":"f888e3d75dc09c1fc8045bf429e699ed468eda1e","Node name for S&R":"easy showAnything"},"widgets_values":["Watermarked:[0, 211, 146, 56, 201]"]},{"id":14,"type":"CSMDialogRun","pos":[359.74139404296875,275.2856140136719],"size":[315,258],"flags":{},"order":3,"mode":0,"inputs":[{"name":"audio0","localized_name":"audio0","type":"AUDIO","shape":7,"link":null},{"name":"audio1","localized_name":"audio1","type":"AUDIO","shape":7,"link":null},{"name":"audio2","localized_name":"audio2","type":"AUDIO","shape":7,"link":null},{"name":"audio3","localized_name":"audio3","type":"AUDIO","shape":7,"link":null},{"name":"text","type":"STRING","widget":{"name":"text"},"link":15}],"outputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","links":[14]}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"CSMDialogRun"},"widgets_values":["","","","","",1]},{"id":15,"type":"MultiLinePromptCSM","pos":[-100.63268280029297,330.7996520996094],"size":[400,200],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"prompt","localized_name":"prompt","type":"STRING","links":[15],"slot_index":0}],"properties":{"aux_id":"billwuhao/ComfyUI_CSM","ver":"8e736d0c2d6a2c9b165b88c048e5110e2c731f59","Node name for S&R":"MultiLinePromptCSM"},"widgets_values":["Hello, what are you doing"]},{"id":9,"type":"PreviewAudio","pos":[732.0776977539062,319.2107849121094],"size":[315,88],"flags":{},"order":6,"mode":0,"inputs":[{"name":"audio","localized_name":"audio","type":"AUDIO","link":14}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.26","Node name for S&R":"PreviewAudio"},"widgets_values":[""]}],"links":[[5,5,0,4,0,"AUDIO"],[6,5,1,3,0,"*"],[9,2,0,5,0,"AUDIO"],[14,14,0,9,0,"AUDIO"],[15,15,0,14,4,"STRING"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.9090909090909091,"offset":[505.47957687377937,121.62480659484862]},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/images/2025-03-18_14-43-49.png b/images/2025-03-18_14-43-49.png new file mode 100644 index 0000000..2fbc4e3 Binary files /dev/null and b/images/2025-03-18_14-43-49.png differ diff --git a/images/2025-03-18_15-38-45.png b/images/2025-03-18_15-38-45.png new file mode 100644 index 0000000..e15c009 Binary files /dev/null and b/images/2025-03-18_15-38-45.png differ diff --git a/images/2025-03-18_19-01-15.png b/images/2025-03-18_19-01-15.png new file mode 100644 index 0000000..624a8ee Binary files /dev/null and b/images/2025-03-18_19-01-15.png differ diff --git a/models.py b/models.py new file mode 100644 index 0000000..22d0927 --- /dev/null +++ b/models.py @@ -0,0 +1,203 @@ +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torchtune +from huggingface_hub import PyTorchModelHubMixin +from torchtune.models import llama3_2 + + +def llama3_2_1B() -> torchtune.modules.transformer.TransformerDecoder: + return llama3_2.llama3_2( + vocab_size=128_256, + num_layers=16, + num_heads=32, + num_kv_heads=8, + embed_dim=2048, + max_seq_len=2048, + intermediate_dim=8192, + attn_dropout=0.0, + norm_eps=1e-5, + rope_base=500_000, + scale_factor=32, + ) + + +def llama3_2_100M() -> torchtune.modules.transformer.TransformerDecoder: + return llama3_2.llama3_2( + vocab_size=128_256, + num_layers=4, + num_heads=8, + num_kv_heads=2, + embed_dim=1024, + max_seq_len=2048, + intermediate_dim=8192, + attn_dropout=0.0, + norm_eps=1e-5, + rope_base=500_000, + scale_factor=32, + ) + + +FLAVORS = { + "llama-1B": llama3_2_1B, + "llama-100M": llama3_2_100M, +} + + +def _prepare_transformer(model): + embed_dim = model.tok_embeddings.embedding_dim + model.tok_embeddings = nn.Identity() + model.output = nn.Identity() + return model, embed_dim + + +def _create_causal_mask(seq_len: int, device: torch.device): + return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)) + + +def _index_causal_mask(mask: torch.Tensor, input_pos: torch.Tensor): + """ + Args: + mask: (max_seq_len, max_seq_len) + input_pos: (batch_size, seq_len) + + Returns: + (batch_size, seq_len, max_seq_len) + """ + r = mask[input_pos, :] + return r + + +def _multinomial_sample_one_no_sync(probs): # Does multinomial sampling without a cuda synchronization + q = torch.empty_like(probs).exponential_(1) + return torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int) + + +def sample_topk(logits: torch.Tensor, topk: int, temperature: float): + logits = logits / temperature + + filter_value: float = -float("Inf") + indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None] + scores_processed = logits.masked_fill(indices_to_remove, filter_value) + scores_processed = torch.nn.functional.log_softmax(scores_processed, dim=-1) + probs = torch.nn.functional.softmax(scores_processed, dim=-1) + + sample_token = _multinomial_sample_one_no_sync(probs) + return sample_token + + +@dataclass +class ModelArgs: + backbone_flavor: str + decoder_flavor: str + text_vocab_size: int + audio_vocab_size: int + audio_num_codebooks: int + + +class Model( + nn.Module, + PyTorchModelHubMixin, + repo_url="https://github.com/SesameAILabs/csm", + pipeline_tag="text-to-speech", + license="apache-2.0", +): + def __init__(self, config: ModelArgs): + super().__init__() + self.config = config + + self.backbone, backbone_dim = _prepare_transformer(FLAVORS[config.backbone_flavor]()) + self.decoder, decoder_dim = _prepare_transformer(FLAVORS[config.decoder_flavor]()) + + self.text_embeddings = nn.Embedding(config.text_vocab_size, backbone_dim) + self.audio_embeddings = nn.Embedding(config.audio_vocab_size * config.audio_num_codebooks, backbone_dim) + + self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False) + self.codebook0_head = nn.Linear(backbone_dim, config.audio_vocab_size, bias=False) + self.audio_head = nn.Parameter(torch.empty(config.audio_num_codebooks - 1, decoder_dim, config.audio_vocab_size)) + + def setup_caches(self, max_batch_size: int) -> torch.Tensor: + """Setup KV caches and return a causal mask.""" + dtype = next(self.parameters()).dtype + device = next(self.parameters()).device + + with device: + self.backbone.setup_caches(max_batch_size, dtype) + self.decoder.setup_caches(max_batch_size, dtype, decoder_max_seq_len=self.config.audio_num_codebooks) + + self.register_buffer("backbone_causal_mask", _create_causal_mask(self.backbone.max_seq_len, device)) + self.register_buffer("decoder_causal_mask", _create_causal_mask(self.config.audio_num_codebooks, device)) + + def generate_frame( + self, + tokens: torch.Tensor, + tokens_mask: torch.Tensor, + input_pos: torch.Tensor, + temperature: float, + topk: int, + ) -> torch.Tensor: + """ + Args: + tokens: (batch_size, seq_len, audio_num_codebooks+1) + tokens_mask: (batch_size, seq_len, audio_num_codebooks+1) + input_pos: (batch_size, seq_len) positions for each token + mask: (batch_size, seq_len, max_seq_len) + + Returns: + (batch_size, audio_num_codebooks) sampled tokens + """ + dtype = next(self.parameters()).dtype + b, s, _ = tokens.size() + + assert self.backbone.caches_are_enabled(), "backbone caches are not enabled" + curr_backbone_mask = _index_causal_mask(self.backbone_causal_mask, input_pos) + embeds = self._embed_tokens(tokens) + masked_embeds = embeds * tokens_mask.unsqueeze(-1) + h = masked_embeds.sum(dim=2) + h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype) + + last_h = h[:, -1, :] + c0_logits = self.codebook0_head(last_h) + c0_sample = sample_topk(c0_logits, topk, temperature) + c0_embed = self._embed_audio(0, c0_sample) + + curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1) + curr_sample = c0_sample.clone() + curr_pos = torch.arange(0, curr_h.size(1), device=curr_h.device).unsqueeze(0).repeat(curr_h.size(0), 1) + + # Decoder caches must be reset every frame. + self.decoder.reset_caches() + for i in range(1, self.config.audio_num_codebooks): + curr_decoder_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos) + decoder_h = self.decoder(self.projection(curr_h), input_pos=curr_pos, mask=curr_decoder_mask).to( + dtype=dtype + ) + ci_logits = torch.mm(decoder_h[:, -1, :], self.audio_head[i - 1]) + ci_sample = sample_topk(ci_logits, topk, temperature) + ci_embed = self._embed_audio(i, ci_sample) + + curr_h = ci_embed + curr_sample = torch.cat([curr_sample, ci_sample], dim=1) + curr_pos = curr_pos[:, -1:] + 1 + + return curr_sample + + def reset_caches(self): + self.backbone.reset_caches() + self.decoder.reset_caches() + + def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor: + return self.audio_embeddings(tokens + codebook * self.config.audio_vocab_size) + + def _embed_tokens(self, tokens: torch.Tensor) -> torch.Tensor: + text_embeds = self.text_embeddings(tokens[:, :, -1]).unsqueeze(-2) + + audio_tokens = tokens[:, :, :-1] + ( + self.config.audio_vocab_size * torch.arange(self.config.audio_num_codebooks, device=tokens.device) + ) + audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape( + tokens.size(0), tokens.size(1), self.config.audio_num_codebooks, -1 + ) + + return torch.cat([audio_embeds, text_embeds], dim=-2) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..87ba43a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "csm_mw" +description = "ComfyUI node of Conversational Speech Model (CSM)." +version = "1.0.0" +license = {file = "LICENSE"} + +[project.urls] +Repository = "https://github.com/billwuhao/ComfyUI_CSM" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "mw" +DisplayName = "ComfyUI_CSM" +Icon = "https://github.com/billwuhao/aiart.website/blob/master/hb.png" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8b640cf --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +torch>=2.4.0 +torchaudio>=2.4.0 +tokenizers>=0.21.0 +transformers>=4.49.0 +huggingface_hub>=0.28.1 +moshi>=0.2.2 +torchtune>=0.4.0 +torchao>=0.9.0 +librosa +silentcipher @ git+https://github.com/SesameAILabs/silentcipher@master \ No newline at end of file