DennisandClaude Opus 4.6
fa5d01b72f
- Fix CFG 1.0 crash: auto-clamp to 1.1 (batch_size mismatch in attention layers) - Extend CFG range from 5.0 to 10.0 for experimental genres - Add comprehensive tag guide to README (genre, voice, mood, structure tags) - Add CFG scale recommendations (2.0 pop, 3.0 rock, 4.0+ electronic) - Document multi-language lyrics support - Add VRAM/stability info from extensive testing - Add .gitignore Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
658 lines
23 KiB
Python
658 lines
23 KiB
Python
"""
|
|
ForgeAI HeartMuLa — ComfyUI Custom Nodes
|
|
All-in-one music generation + lyrics transcription for 16 GB GPUs.
|
|
|
|
Fixes included:
|
|
- transformers 5.x compatibility (ignore_mismatched_sizes)
|
|
- torchtune >= 0.5 RoPE init fix
|
|
- OOM fix: model.cpu() offload before codec decode
|
|
- soundfile instead of torchaudio (no torchcodec needed)
|
|
- bitsandbytes 4-bit/8-bit quantization support
|
|
- Pre-quantized 4-bit checkpoint auto-detection
|
|
- Fade-out post-processing
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import torch
|
|
import numpy as np
|
|
import folder_paths
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model path setup
|
|
# ---------------------------------------------------------------------------
|
|
HEARTMULA_DIR = os.path.join(folder_paths.models_dir, "HeartMuLa")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scan for available HeartMuLa models
|
|
# ---------------------------------------------------------------------------
|
|
def _scan_model_folders():
|
|
"""Scan HeartMuLa directory for available model folders."""
|
|
if not os.path.exists(HEARTMULA_DIR):
|
|
return ["HeartMuLa-oss-3B"]
|
|
|
|
folders = []
|
|
skip = {"HeartCodec-oss", "heartlib", "HeartMuLa-oss-whisper"}
|
|
for name in sorted(os.listdir(HEARTMULA_DIR)):
|
|
path = os.path.join(HEARTMULA_DIR, name)
|
|
if os.path.isdir(path) and name not in skip:
|
|
# Check if it looks like a model folder (has config.json)
|
|
if os.path.exists(os.path.join(path, "config.json")):
|
|
folders.append(name)
|
|
|
|
return folders if folders else ["HeartMuLa-oss-3B"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Patch heartlib on import: RoPE init fix for torchtune >= 0.5
|
|
# ---------------------------------------------------------------------------
|
|
def _apply_heartlib_patches():
|
|
"""Monkey-patch HeartMuLa.setup_caches to call rope_init() for torchtune >= 0.5"""
|
|
try:
|
|
from heartlib.heartmula.modeling_heartmula import HeartMuLa
|
|
|
|
_original_setup_caches = HeartMuLa.setup_caches
|
|
|
|
def _patched_setup_caches(self, batch_size, *args, **kwargs):
|
|
_original_setup_caches(self, batch_size, *args, **kwargs)
|
|
# RoPE init fix: torchtune >= 0.5 requires explicit rope_init()
|
|
for m in self.modules():
|
|
if hasattr(m, "rope_init"):
|
|
m.rope_init()
|
|
# Move RoPE buffers to same device as model
|
|
try:
|
|
p = next(self.parameters())
|
|
m.to(p.device)
|
|
except StopIteration:
|
|
pass
|
|
|
|
HeartMuLa.setup_caches = _patched_setup_caches
|
|
print("[ForgeAI] HeartMuLa RoPE patch applied")
|
|
except ImportError:
|
|
print("[ForgeAI] WARNING: heartlib not found, install with: pip install heartlib")
|
|
|
|
|
|
_apply_heartlib_patches()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline cache (keep models loaded between runs)
|
|
# ---------------------------------------------------------------------------
|
|
_pipeline_cache = {}
|
|
|
|
|
|
def _get_device():
|
|
return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
def _load_pipeline(model_folder, quantize="none"):
|
|
"""Load HeartMuLa generation pipeline with optional quantization.
|
|
|
|
model_folder: name of the folder inside HeartMuLa/ (e.g. HeartMuLa-4bit-3B)
|
|
"""
|
|
cache_key = f"{model_folder}_{quantize}"
|
|
if cache_key in _pipeline_cache:
|
|
return _pipeline_cache[cache_key]
|
|
|
|
from tokenizers import Tokenizer
|
|
from heartlib.heartmula.modeling_heartmula import HeartMuLa
|
|
from heartlib.heartcodec.modeling_heartcodec import HeartCodec
|
|
|
|
device = _get_device()
|
|
dtype = torch.bfloat16
|
|
|
|
# --- Load HeartCodec ---
|
|
codec_path = os.path.join(HEARTMULA_DIR, "HeartCodec-oss")
|
|
if not os.path.exists(codec_path):
|
|
raise FileNotFoundError(
|
|
f"HeartCodec not found at {codec_path}. "
|
|
f"Download from https://huggingface.co/HeartMuLa/HeartMuLa-oss-3B"
|
|
)
|
|
print(f"[ForgeAI] Loading HeartCodec from {codec_path}")
|
|
codec = HeartCodec.from_pretrained(
|
|
codec_path, device_map=device, ignore_mismatched_sizes=True
|
|
)
|
|
|
|
# --- Determine model path and quantization config ---
|
|
model_path = os.path.join(HEARTMULA_DIR, model_folder)
|
|
bnb_config = None
|
|
device_map = None
|
|
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(
|
|
f"Model not found at {model_path}. "
|
|
f"Download from https://huggingface.co/HeartMuLa/HeartMuLa-oss-3B "
|
|
f"or https://huggingface.co/PavonicAI/HeartMuLa-3B-4bit"
|
|
)
|
|
|
|
# Check if the model folder has a baked-in quantization config
|
|
import json
|
|
config_path = os.path.join(model_path, "config.json")
|
|
has_baked_quant = False
|
|
if os.path.exists(config_path):
|
|
with open(config_path, encoding="utf-8") as f:
|
|
model_cfg = json.load(f)
|
|
if "quantization_config" in model_cfg:
|
|
has_baked_quant = True
|
|
print(f"[ForgeAI] Model has baked-in quantization config")
|
|
|
|
if has_baked_quant:
|
|
# Pre-quantized checkpoint — just load it, config.json has everything
|
|
device_map = "cuda:0"
|
|
print(f"[ForgeAI] Loading pre-quantized checkpoint: {model_folder}")
|
|
elif quantize == "4bit":
|
|
from transformers import BitsAndBytesConfig
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
bnb_4bit_quant_type="nf4",
|
|
)
|
|
device_map = "cuda:0"
|
|
print(f"[ForgeAI] Quantizing on-the-fly (4bit): {model_folder}")
|
|
elif quantize == "8bit":
|
|
from transformers import BitsAndBytesConfig
|
|
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
|
device_map = "cuda:0"
|
|
print(f"[ForgeAI] Quantizing on-the-fly (8bit): {model_folder}")
|
|
else:
|
|
print(f"[ForgeAI] Loading full precision: {model_folder}")
|
|
|
|
print(f"[ForgeAI] Loading {model_folder} ({quantize})...")
|
|
model = HeartMuLa.from_pretrained(
|
|
model_path,
|
|
dtype=dtype,
|
|
quantization_config=bnb_config,
|
|
device_map=device_map,
|
|
ignore_mismatched_sizes=True,
|
|
)
|
|
|
|
# --- Load tokenizer + config ---
|
|
tokenizer_path = os.path.join(HEARTMULA_DIR, "tokenizer.json")
|
|
if not os.path.exists(tokenizer_path):
|
|
raise FileNotFoundError(f"tokenizer.json not found at {tokenizer_path}")
|
|
tokenizer = Tokenizer.from_file(tokenizer_path)
|
|
|
|
gen_config_path = os.path.join(HEARTMULA_DIR, "gen_config.json")
|
|
if not os.path.exists(gen_config_path):
|
|
raise FileNotFoundError(f"gen_config.json not found at {gen_config_path}")
|
|
|
|
with open(gen_config_path, encoding="utf-8") as f:
|
|
gen_cfg = json.load(f)
|
|
|
|
pipeline_data = {
|
|
"model": model,
|
|
"codec": codec,
|
|
"tokenizer": tokenizer,
|
|
"gen_config": gen_cfg,
|
|
"device": device,
|
|
"dtype": dtype,
|
|
"num_quantizers": codec.config.num_quantizers,
|
|
"muq_dim": model.config.muq_dim,
|
|
}
|
|
|
|
_pipeline_cache[cache_key] = pipeline_data
|
|
print(f"[ForgeAI] Pipeline ready ({quantize})")
|
|
return pipeline_data
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generation logic (all fixes baked in)
|
|
# ---------------------------------------------------------------------------
|
|
def _generate_music(
|
|
pipeline_data,
|
|
tags: str,
|
|
lyrics: str,
|
|
max_duration_sec: int = 30,
|
|
temperature: float = 1.0,
|
|
top_k: int = 50,
|
|
cfg_scale: float = 1.5,
|
|
):
|
|
"""Generate music from lyrics. Returns (waveform_tensor, sample_rate)."""
|
|
from comfy.utils import ProgressBar
|
|
|
|
model = pipeline_data["model"]
|
|
codec = pipeline_data["codec"]
|
|
tokenizer = pipeline_data["tokenizer"]
|
|
gen_cfg = pipeline_data["gen_config"]
|
|
device = pipeline_data["device"]
|
|
dtype = pipeline_data["dtype"]
|
|
parallel_number = pipeline_data["num_quantizers"] + 1
|
|
muq_dim = pipeline_data["muq_dim"]
|
|
|
|
text_bos_id = gen_cfg.get("text_bos_id", 128000)
|
|
text_eos_id = gen_cfg.get("text_eos_id", 128001)
|
|
audio_eos_id = gen_cfg.get("audio_eos_id", 8193)
|
|
empty_id = gen_cfg.get("empty_id", 0)
|
|
|
|
# --- Preprocess tags ---
|
|
tags = tags.lower().strip()
|
|
if not tags.startswith("<tag>"):
|
|
tags = f"<tag>{tags}"
|
|
if not tags.endswith("</tag>"):
|
|
tags = f"{tags}</tag>"
|
|
|
|
tags_ids = tokenizer.encode(tags).ids
|
|
if tags_ids[0] != text_bos_id:
|
|
tags_ids = [text_bos_id] + tags_ids
|
|
if tags_ids[-1] != text_eos_id:
|
|
tags_ids = tags_ids + [text_eos_id]
|
|
|
|
# --- Preprocess lyrics ---
|
|
lyrics = lyrics.lower().strip()
|
|
lyrics_ids = tokenizer.encode(lyrics).ids
|
|
if lyrics_ids[0] != text_bos_id:
|
|
lyrics_ids = [text_bos_id] + lyrics_ids
|
|
if lyrics_ids[-1] != text_eos_id:
|
|
lyrics_ids = lyrics_ids + [text_eos_id]
|
|
|
|
# --- Build prompt tokens ---
|
|
muq_embed = torch.zeros([muq_dim], dtype=dtype)
|
|
muq_idx = len(tags_ids)
|
|
prompt_len = len(tags_ids) + 1 + len(lyrics_ids)
|
|
|
|
tokens = torch.zeros([prompt_len, parallel_number], dtype=torch.long)
|
|
tokens[: len(tags_ids), -1] = torch.tensor(tags_ids)
|
|
tokens[len(tags_ids) + 1 :, -1] = torch.tensor(lyrics_ids)
|
|
|
|
tokens_mask = torch.zeros_like(tokens, dtype=torch.bool)
|
|
tokens_mask[:, -1] = True
|
|
|
|
# CFG 1.0 causes batch_size mismatch in attention layers
|
|
# Clamp to minimum 1.1 to always use batch_size=2
|
|
if cfg_scale <= 1.0:
|
|
print(f"[ForgeAI] CFG {cfg_scale} clamped to 1.1 (minimum for stable generation)")
|
|
cfg_scale = 1.1
|
|
bs_size = 2
|
|
|
|
def _cfg_cat(tensor, cfg_s):
|
|
tensor = tensor.unsqueeze(0)
|
|
if cfg_s != 1.0:
|
|
tensor = torch.cat([tensor, tensor], dim=0)
|
|
return tensor
|
|
|
|
prompt_tokens = _cfg_cat(tokens, cfg_scale).to(device)
|
|
prompt_tokens_mask = _cfg_cat(tokens_mask, cfg_scale).to(device)
|
|
continuous_segment = _cfg_cat(muq_embed, cfg_scale).to(device)
|
|
starts = [muq_idx] * bs_size
|
|
prompt_pos = _cfg_cat(
|
|
torch.arange(prompt_len, dtype=torch.long), cfg_scale
|
|
).to(device)
|
|
|
|
# --- Generate frames with ComfyUI progress bar ---
|
|
max_audio_frames = (max_duration_sec * 1000) // 80
|
|
frames = []
|
|
|
|
model.setup_caches(bs_size)
|
|
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
curr_token = model.generate_frame(
|
|
tokens=prompt_tokens,
|
|
tokens_mask=prompt_tokens_mask,
|
|
input_pos=prompt_pos,
|
|
temperature=temperature,
|
|
topk=top_k,
|
|
cfg_scale=cfg_scale,
|
|
continuous_segments=continuous_segment,
|
|
starts=starts,
|
|
)
|
|
frames.append(curr_token[0:1])
|
|
|
|
print(f"[ForgeAI] Generating {max_duration_sec}s of audio...")
|
|
pbar = ProgressBar(max_audio_frames)
|
|
for i in range(max_audio_frames):
|
|
padded_token = (
|
|
torch.ones(
|
|
(curr_token.shape[0], parallel_number),
|
|
device=curr_token.device,
|
|
dtype=torch.long,
|
|
)
|
|
* empty_id
|
|
)
|
|
padded_token[:, :-1] = curr_token
|
|
padded_token = padded_token.unsqueeze(1)
|
|
padded_token_mask = torch.ones_like(padded_token, dtype=torch.bool)
|
|
padded_token_mask[..., -1] = False
|
|
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
curr_token = model.generate_frame(
|
|
tokens=padded_token,
|
|
tokens_mask=padded_token_mask,
|
|
input_pos=prompt_pos[..., -1:] + i + 1,
|
|
temperature=temperature,
|
|
topk=top_k,
|
|
cfg_scale=cfg_scale,
|
|
continuous_segments=None,
|
|
starts=None,
|
|
)
|
|
if torch.any(curr_token[0:1, :] >= audio_eos_id):
|
|
print(f"[ForgeAI] Audio EOS reached at frame {i}")
|
|
break
|
|
frames.append(curr_token[0:1])
|
|
pbar.update(1)
|
|
|
|
# --- Decode audio (with OOM fix: offload model first) ---
|
|
frames_tensor = torch.stack(frames).permute(1, 2, 0).squeeze(0)
|
|
model.reset_caches()
|
|
model.cpu()
|
|
torch.cuda.empty_cache()
|
|
print("[ForgeAI] Model offloaded, decoding audio...")
|
|
|
|
wav = codec.detokenize(frames_tensor)
|
|
|
|
# Move model back to GPU for next run
|
|
try:
|
|
model.to(device)
|
|
except Exception:
|
|
pass # quantized models may not support .to()
|
|
|
|
return wav, 48000
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Post-processing: fade out
|
|
# ---------------------------------------------------------------------------
|
|
def _apply_fade_out(wav_np, sample_rate, fade_seconds=3.0):
|
|
"""Apply a smooth fade-out to the last N seconds of audio."""
|
|
fade_samples = int(fade_seconds * sample_rate)
|
|
if fade_samples >= len(wav_np):
|
|
fade_samples = len(wav_np)
|
|
|
|
# Create linear fade curve
|
|
fade_curve = np.linspace(1.0, 0.0, fade_samples)
|
|
|
|
if wav_np.ndim == 2:
|
|
# Multi-channel: apply to all channels
|
|
wav_np[-fade_samples:] *= fade_curve[:, np.newaxis]
|
|
else:
|
|
wav_np[-fade_samples:] *= fade_curve
|
|
|
|
return wav_np
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Save audio to file (WAV or MP3) — uses PyAV for MP3 (no ffmpeg needed)
|
|
# ---------------------------------------------------------------------------
|
|
def _save_audio(wav_tensor, sample_rate, save_path, output_format="wav",
|
|
fade_out_seconds=0.0):
|
|
"""Save audio tensor to file. Uses PyAV for MP3 encoding."""
|
|
import soundfile as sf
|
|
|
|
wav_np = wav_tensor.cpu().float().numpy()
|
|
if wav_np.ndim == 2:
|
|
wav_np = wav_np.T # soundfile expects (samples, channels)
|
|
|
|
# Apply fade-out if requested
|
|
if fade_out_seconds > 0:
|
|
wav_np = _apply_fade_out(wav_np, sample_rate, fade_out_seconds)
|
|
print(f"[ForgeAI] Applied {fade_out_seconds}s fade-out")
|
|
|
|
if output_format == "mp3":
|
|
try:
|
|
import av
|
|
|
|
# Encode MP3 using PyAV (bundled with ComfyUI, no external ffmpeg needed)
|
|
container = av.open(save_path, mode="w")
|
|
channels = wav_np.shape[1] if wav_np.ndim == 2 else 1
|
|
stream = container.add_stream("libmp3lame", rate=sample_rate)
|
|
stream.bit_rate = 320000 # 320 kbps
|
|
stream.layout = "stereo" if channels > 1 else "mono"
|
|
|
|
# PyAV expects (samples, channels) float32
|
|
if wav_np.ndim == 1:
|
|
wav_np = wav_np.reshape(-1, 1)
|
|
|
|
# Process in chunks to avoid memory issues
|
|
chunk_size = sample_rate # 1 second chunks
|
|
for start in range(0, len(wav_np), chunk_size):
|
|
chunk = wav_np[start : start + chunk_size]
|
|
frame = av.AudioFrame.from_ndarray(
|
|
chunk.T, # PyAV expects (channels, samples)
|
|
format="fltp",
|
|
layout=stream.layout.name,
|
|
)
|
|
frame.sample_rate = sample_rate
|
|
for packet in stream.encode(frame):
|
|
container.mux(packet)
|
|
|
|
# Flush encoder
|
|
for packet in stream.encode(None):
|
|
container.mux(packet)
|
|
container.close()
|
|
print(f"[ForgeAI] Saved MP3: {save_path}")
|
|
|
|
except Exception as e:
|
|
# Fallback to WAV if MP3 encoding fails
|
|
print(f"[ForgeAI] MP3 encoding failed ({e}), saving as WAV")
|
|
save_path = save_path.rsplit(".", 1)[0] + ".wav"
|
|
sf.write(save_path, wav_np, sample_rate)
|
|
print(f"[ForgeAI] Saved WAV: {save_path}")
|
|
else:
|
|
sf.write(save_path, wav_np, sample_rate)
|
|
print(f"[ForgeAI] Saved WAV: {save_path}")
|
|
|
|
return save_path
|
|
|
|
|
|
# ===========================================================================
|
|
# ComfyUI Node: ForgeAI Music Generator
|
|
# ===========================================================================
|
|
class ForgeAI_HeartMuLa_Generate:
|
|
"""
|
|
All-in-one HeartMuLa music generation node.
|
|
Generates music from lyrics + genre tags. Supports 4-bit quantization
|
|
for 16 GB GPUs. Saves output as WAV or MP3.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_folders = _scan_model_folders()
|
|
return {
|
|
"required": {
|
|
"lyrics": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "[verse]\nHello world\nThis is a song\n\n[chorus]\nLa la la",
|
|
},
|
|
),
|
|
"tags": (
|
|
"STRING",
|
|
{
|
|
"default": "warm, Piano, pop, acoustic guitar",
|
|
},
|
|
),
|
|
"model": (model_folders,),
|
|
"quantize": (["4bit", "8bit", "none"],),
|
|
"duration_seconds": (
|
|
"INT",
|
|
{"default": 120, "min": 5, "max": 240, "step": 5},
|
|
),
|
|
"cfg_scale": (
|
|
"FLOAT",
|
|
{"default": 2.0, "min": 1.0, "max": 10.0, "step": 0.1},
|
|
),
|
|
"temperature": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.1, "max": 2.0, "step": 0.05},
|
|
),
|
|
"top_k": (
|
|
"INT",
|
|
{"default": 50, "min": 1, "max": 500, "step": 1},
|
|
),
|
|
"fade_out": (
|
|
"FLOAT",
|
|
{"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.5},
|
|
),
|
|
"output_format": (["wav", "mp3"],),
|
|
"filename_prefix": (
|
|
"STRING",
|
|
{"default": "ForgeAI_music"},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO", "STRING")
|
|
RETURN_NAMES = ("audio", "file_path")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "ForgeAI/Audio"
|
|
OUTPUT_NODE = True
|
|
|
|
def generate(
|
|
self,
|
|
lyrics,
|
|
tags,
|
|
model,
|
|
quantize,
|
|
duration_seconds,
|
|
cfg_scale,
|
|
temperature,
|
|
top_k,
|
|
fade_out,
|
|
output_format,
|
|
filename_prefix,
|
|
):
|
|
# Load pipeline (model is now folder name directly)
|
|
pipeline = _load_pipeline(model, quantize)
|
|
|
|
# Generate audio
|
|
wav, sample_rate = _generate_music(
|
|
pipeline,
|
|
tags=tags,
|
|
lyrics=lyrics,
|
|
max_duration_sec=duration_seconds,
|
|
temperature=temperature,
|
|
top_k=top_k,
|
|
cfg_scale=cfg_scale,
|
|
)
|
|
|
|
# Build output path
|
|
import time
|
|
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
|
ext = output_format
|
|
filename = f"{filename_prefix}_{timestamp}.{ext}"
|
|
save_path = os.path.join(self.output_dir, filename)
|
|
|
|
# Save file (with optional fade-out)
|
|
actual_path = _save_audio(wav, sample_rate, save_path, output_format,
|
|
fade_out_seconds=fade_out)
|
|
actual_filename = os.path.basename(actual_path)
|
|
|
|
# Build ComfyUI AUDIO output
|
|
wav_np = wav.cpu().float().numpy()
|
|
waveform = torch.from_numpy(wav_np)
|
|
if waveform.ndim == 1:
|
|
waveform = waveform.unsqueeze(0)
|
|
elif waveform.ndim == 2 and waveform.shape[0] > waveform.shape[1]:
|
|
waveform = waveform.T
|
|
audio_output = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
|
|
|
|
# Return with UI dict so ComfyUI shows audio preview in assets
|
|
return {
|
|
"ui": {
|
|
"audio": [
|
|
{
|
|
"filename": actual_filename,
|
|
"subfolder": "",
|
|
"type": "output",
|
|
}
|
|
]
|
|
},
|
|
"result": (audio_output, actual_path),
|
|
}
|
|
|
|
|
|
# ===========================================================================
|
|
# ComfyUI Node: ForgeAI Lyrics Transcriber
|
|
# ===========================================================================
|
|
class ForgeAI_HeartMuLa_Transcribe:
|
|
"""
|
|
Transcribe lyrics from audio using HeartMuLa's Whisper-based transcriber.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("lyrics",)
|
|
FUNCTION = "transcribe"
|
|
CATEGORY = "ForgeAI/Audio"
|
|
|
|
def transcribe(self, audio):
|
|
import soundfile as sf
|
|
import tempfile
|
|
|
|
# Extract waveform from ComfyUI AUDIO type
|
|
waveform = audio["waveform"].squeeze(0) # remove batch dim
|
|
sample_rate = audio["sample_rate"]
|
|
|
|
# Save to temp file for Whisper
|
|
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
|
|
tmp_path = tmp.name
|
|
wav_np = waveform.cpu().float().numpy()
|
|
if wav_np.ndim == 2:
|
|
wav_np = wav_np.T
|
|
sf.write(tmp_path, wav_np, sample_rate)
|
|
|
|
try:
|
|
# Load Whisper model
|
|
from transformers import WhisperProcessor, WhisperForConditionalGeneration
|
|
import librosa
|
|
|
|
whisper_path = os.path.join(HEARTMULA_DIR, "HeartMuLa-oss-whisper")
|
|
if not os.path.exists(whisper_path):
|
|
# Fall back to standard whisper
|
|
whisper_path = "openai/whisper-large-v3"
|
|
print(f"[ForgeAI] HeartMuLa whisper not found, using {whisper_path}")
|
|
|
|
processor = WhisperProcessor.from_pretrained(whisper_path)
|
|
whisper_model = WhisperForConditionalGeneration.from_pretrained(
|
|
whisper_path,
|
|
ignore_mismatched_sizes=True,
|
|
).to(_get_device())
|
|
|
|
# Load and resample audio to 16kHz for Whisper
|
|
audio_data, sr = librosa.load(tmp_path, sr=16000)
|
|
input_features = processor(
|
|
audio_data, sampling_rate=16000, return_tensors="pt"
|
|
).input_features.to(_get_device())
|
|
|
|
# Transcribe
|
|
predicted_ids = whisper_model.generate(input_features)
|
|
transcription = processor.batch_decode(
|
|
predicted_ids, skip_special_tokens=True
|
|
)[0]
|
|
|
|
# Clean up whisper model
|
|
del whisper_model
|
|
torch.cuda.empty_cache()
|
|
|
|
print(f"[ForgeAI] Transcribed: {transcription[:100]}...")
|
|
return (transcription,)
|
|
|
|
finally:
|
|
os.unlink(tmp_path)
|
|
|
|
|
|
# ===========================================================================
|
|
# ComfyUI Registration
|
|
# ===========================================================================
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ForgeAI_HeartMuLa_Generate": ForgeAI_HeartMuLa_Generate,
|
|
"ForgeAI_HeartMuLa_Transcribe": ForgeAI_HeartMuLa_Transcribe,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ForgeAI_HeartMuLa_Generate": "ForgeAI Music Generator (HeartMuLa)",
|
|
"ForgeAI_HeartMuLa_Transcribe": "ForgeAI Lyrics Transcriber",
|
|
}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|