Delete nodes directory
This commit is contained in:
Shmuel Ronen
@@ -1,12 +0,0 @@
|
||||
from .model_loader import ChatterboxModelLoader
|
||||
from .generate import ChatterboxGenerate
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ChatterboxModelLoader": ChatterboxModelLoader,
|
||||
"ChatterboxGenerate": ChatterboxGenerate,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ChatterboxModelLoader": "Chatterbox Model Loader",
|
||||
"ChatterboxGenerate": "Chatterbox Generate",
|
||||
}
|
||||
@@ -1,153 +0,0 @@
|
||||
import torch
|
||||
import torchaudio
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
OUTPUT_DIR = os.path.join(folder_paths.output_directory, "audio")
|
||||
except ImportError:
|
||||
OUTPUT_DIR = str(Path.home() / "ComfyUI" / "output" / "audio")
|
||||
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
class ChatterboxGenerate:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("CHATTERBOX_MODEL",),
|
||||
"text": ("STRING", {"multiline": True, "default": "Hello, this is a test of the Chatterbox TTS system."}),
|
||||
"exaggeration": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.1}),
|
||||
"cfg_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}),
|
||||
"temperature": ("FLOAT", {"default": 0.55, "min": 0.1, "max": 2.0, "step": 0.1}),
|
||||
},
|
||||
"optional": {
|
||||
"reference_audio": ("AUDIO",),
|
||||
"filename": ("STRING", {"default": "chatterbox_output"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AUDIO",)
|
||||
RETURN_NAMES = ("audio",)
|
||||
FUNCTION = "generate_speech"
|
||||
CATEGORY = "ChatterBox"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def generate_speech(self, model, text, exaggeration=0.5, cfg_weight=0.5, temperature=0.8, reference_audio=None, filename="chatterbox_output"):
|
||||
audio_tensor = None # Initialize to avoid UnboundLocalError
|
||||
try:
|
||||
# Prepare audio prompt if provided
|
||||
audio_prompt_path = None
|
||||
if reference_audio is not None:
|
||||
print(f"Reference audio type: {type(reference_audio)}")
|
||||
print(f"Reference audio content: {reference_audio if isinstance(reference_audio, dict) else 'tensor'}")
|
||||
|
||||
# Handle ComfyUI audio format (usually a dict with 'waveform' and 'sample_rate')
|
||||
if isinstance(reference_audio, dict):
|
||||
if 'waveform' in reference_audio:
|
||||
audio_data = reference_audio['waveform']
|
||||
sample_rate = reference_audio.get('sample_rate', model.sr)
|
||||
else:
|
||||
print(f"Warning: Unknown audio dict format: {reference_audio.keys()}")
|
||||
audio_data = None
|
||||
elif torch.is_tensor(reference_audio):
|
||||
audio_data = reference_audio
|
||||
sample_rate = model.sr
|
||||
else:
|
||||
print(f"Warning: Unknown audio format: {type(reference_audio)}")
|
||||
audio_data = None
|
||||
|
||||
if audio_data is not None:
|
||||
# Save audio prompt to temporary file
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
|
||||
audio_prompt_path = tmp_file.name
|
||||
|
||||
# Ensure audio_data is 2D for torchaudio.save
|
||||
if audio_data.dim() == 1:
|
||||
audio_data = audio_data.unsqueeze(0)
|
||||
elif audio_data.dim() > 2:
|
||||
audio_data = audio_data.squeeze()
|
||||
if audio_data.dim() == 1:
|
||||
audio_data = audio_data.unsqueeze(0)
|
||||
|
||||
torchaudio.save(audio_prompt_path, audio_data.float(), sample_rate)
|
||||
print(f"Saved reference audio to: {audio_prompt_path}")
|
||||
|
||||
# Generate speech
|
||||
print(f"Generating speech: '{text[:50]}{'...' if len(text) > 50 else ''}'")
|
||||
audio_tensor = model.generate(
|
||||
text=text,
|
||||
audio_prompt_path=audio_prompt_path,
|
||||
exaggeration=exaggeration,
|
||||
cfg_weight=cfg_weight,
|
||||
temperature=temperature
|
||||
)
|
||||
|
||||
# Fix tensor shape: squeeze to remove extra dimensions and ensure 2D (channels, samples)
|
||||
print(f"Original audio tensor shape: {audio_tensor.shape}")
|
||||
|
||||
# Remove any extra dimensions and ensure we have (channels, samples)
|
||||
audio_tensor = audio_tensor.squeeze() # Remove singleton dimensions
|
||||
|
||||
# If 1D, add channel dimension to make it (1, samples)
|
||||
if audio_tensor.dim() == 1:
|
||||
audio_tensor = audio_tensor.unsqueeze(0)
|
||||
# If still too many dimensions, take the last 2
|
||||
elif audio_tensor.dim() > 2:
|
||||
audio_tensor = audio_tensor.view(-1, audio_tensor.shape[-1])
|
||||
if audio_tensor.shape[0] > 1:
|
||||
audio_tensor = audio_tensor[0:1] # Take only first channel
|
||||
|
||||
print(f"Fixed audio tensor shape: {audio_tensor.shape}")
|
||||
|
||||
# Save to output directory
|
||||
if not filename.endswith('.wav'):
|
||||
filename += '.wav'
|
||||
|
||||
output_path = os.path.join(OUTPUT_DIR, filename)
|
||||
|
||||
# Ensure tensor is in correct format for saving (2D: channels, samples)
|
||||
save_tensor = audio_tensor.squeeze(0) if audio_tensor.dim() == 3 else audio_tensor
|
||||
save_tensor = save_tensor.float() # Ensure float type
|
||||
|
||||
torchaudio.save(
|
||||
output_path,
|
||||
save_tensor,
|
||||
model.sr
|
||||
)
|
||||
|
||||
print(f"Audio saved to: {output_path}")
|
||||
|
||||
# Clean up temporary audio prompt file
|
||||
if audio_prompt_path and os.path.exists(audio_prompt_path):
|
||||
os.unlink(audio_prompt_path)
|
||||
|
||||
# Return audio in ComfyUI format: {"waveform": tensor, "sample_rate": int}
|
||||
# Ensure tensor has batch dimension: [batch, channels, samples]
|
||||
if audio_tensor.dim() == 2: # [channels, samples]
|
||||
audio_tensor = audio_tensor.unsqueeze(0) # Add batch dimension -> [1, channels, samples]
|
||||
|
||||
audio_output = {
|
||||
"waveform": audio_tensor,
|
||||
"sample_rate": model.sr
|
||||
}
|
||||
|
||||
print(f"Audio output format: waveform shape={audio_output['waveform'].shape}, sample_rate={audio_output['sample_rate']}")
|
||||
|
||||
return (audio_output,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in speech generation: {e}")
|
||||
if audio_tensor is not None:
|
||||
print(f"Audio tensor info: shape={getattr(audio_tensor, 'shape', 'N/A')}, dtype={getattr(audio_tensor, 'dtype', 'N/A')}")
|
||||
else:
|
||||
print("Audio tensor was not created")
|
||||
|
||||
# Clean up temporary audio prompt file if it exists
|
||||
if 'audio_prompt_path' in locals() and audio_prompt_path and os.path.exists(audio_prompt_path):
|
||||
os.unlink(audio_prompt_path)
|
||||
|
||||
raise e
|
||||
@@ -1,82 +0,0 @@
|
||||
import torch
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from ..chatterbox.tts import ChatterboxTTS
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
MODELS_DIR = folder_paths.models_dir
|
||||
except ImportError:
|
||||
MODELS_DIR = str(Path.home() / "ComfyUI" / "models")
|
||||
|
||||
CHATTERBOX_MODEL_DIR = os.path.join(MODELS_DIR, "TTS", "chatterbox")
|
||||
os.makedirs(CHATTERBOX_MODEL_DIR, exist_ok=True)
|
||||
|
||||
_MODEL_CACHE = {}
|
||||
|
||||
class ChatterboxModelLoader:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"device": (["auto", "cuda", "cpu", "mps"], {"default": "auto"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CHATTERBOX_MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "ChatterBox"
|
||||
|
||||
def load_model(self, device):
|
||||
if device == "auto":
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
elif torch.backends.mps.is_available():
|
||||
device = "mps"
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
cache_key = f"chatterbox_{device}"
|
||||
|
||||
if cache_key in _MODEL_CACHE:
|
||||
print(f"Using cached model on {device}")
|
||||
return (_MODEL_CACHE[cache_key],)
|
||||
|
||||
# Check if manual model files exist
|
||||
model_path = Path(CHATTERBOX_MODEL_DIR)
|
||||
required_files = ["ve.safetensors", "t3_cfg.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"]
|
||||
|
||||
missing_files = []
|
||||
for file in required_files:
|
||||
if not (model_path / file).exists():
|
||||
missing_files.append(file)
|
||||
|
||||
if missing_files:
|
||||
error_msg = f"Missing model files in {CHATTERBOX_MODEL_DIR}:\n"
|
||||
for file in missing_files:
|
||||
error_msg += f" - {file}\n"
|
||||
error_msg += "\nPlease download these files manually from:\n"
|
||||
error_msg += "https://huggingface.co/ResembleAI/chatterbox/tree/main\n"
|
||||
error_msg += f"and place them in: {CHATTERBOX_MODEL_DIR}"
|
||||
|
||||
print(error_msg)
|
||||
raise FileNotFoundError(error_msg)
|
||||
|
||||
try:
|
||||
print(f"Loading model from {CHATTERBOX_MODEL_DIR}")
|
||||
print("All required files found, loading model...")
|
||||
|
||||
model = ChatterboxTTS.from_local(model_path, device=device)
|
||||
|
||||
_MODEL_CACHE[cache_key] = model
|
||||
print(f"Model loaded successfully on {device}")
|
||||
|
||||
return (model,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error loading model: {e}")
|
||||
print("Make sure all model files are properly downloaded and not corrupted.")
|
||||
raise e
|
||||
Reference in New Issue
Block a user