Spaces:
Runtime error
Runtime error
| import type { GenerationRequest } from "@/lib/types"; | |
| export type ComfyNode = { | |
| class_type: string; | |
| inputs: Record<string, unknown>; | |
| _meta?: { title: string }; | |
| }; | |
| export type ComfyPrompt = Record<string, ComfyNode>; | |
| /** | |
| * The community Real-Audio decoder adapter. Its author reports v9 as the best | |
| * of the published ladder, so that is what this ships. It attaches to the | |
| * MODEL side only and was built against the bf16 layout, both of which the | |
| * graph below respects. | |
| * | |
| * The matching tokenizer head, which is what would actually read a recording | |
| * into real-audio tokens, has no ComfyUI node and runs only from the repo's | |
| * python scripts. So this adapter shifts how the renderer sounds; it does not | |
| * give the model a second channel into the reference. | |
| */ | |
| export const REAL_AUDIO_LORA = "nar_lora_joint_v9_comfyui.safetensors"; | |
| export const REAL_AUDIO_LORAS = [REAL_AUDIO_LORA]; | |
| export const REAL_AUDIO_BASE_CHECKPOINT = "yue2_3b_bf16.safetensors"; | |
| /** | |
| * The instrumental chain-of-thought adapter, trained on ~2.7k instrumental | |
| * tracks paired with chord-annotated ABC from SheetSage2. It steers the | |
| * autoregressive side, so it rides on CLIP, and its training assumed a full | |
| * score was attached: its author's note is `cot="full"` with a non-empty ABC | |
| * step, which is what INSTRUMENTAL_REQUIRES below encodes. | |
| * | |
| * Because every one of those training tracks was instrumental, it pulls away | |
| * from vocals. That is a feature when the goal is an arrangement and a | |
| * problem when there are lyrics to sing, so the UI says so rather than | |
| * letting it quietly mute a vocal take. | |
| */ | |
| export const INSTRUMENTAL_LORA = "ar_lora_inst_v3abc_comfyui.safetensors"; | |
| export const INSTRUMENTAL_LORAS = [INSTRUMENTAL_LORA]; | |
| /** Whether the graph can honour the instrumental adapter as trained. */ | |
| export function instrumentalRequirementsMet(req: { | |
| mode: GenerationRequest["mode"]; | |
| scoreSource: GenerationRequest["scoreSource"]; | |
| }): boolean { | |
| return req.mode === "full" && req.scoreSource !== "none"; | |
| } | |
| /** Node ids are stable so the websocket relay can map progress onto a stage. */ | |
| export const NODE = { | |
| checkpoint: "ckpt", | |
| audioEncoder: "audio_encoder", | |
| loadAudio: "load_audio", | |
| trim: "trim", | |
| transcribe: "transcribe", | |
| abcRelay: "abc_relay", | |
| loraNar: "lora_nar", | |
| loraInst: "lora_inst", | |
| planner: "planner", | |
| plannerRelay: "planner_relay", | |
| music: "music", | |
| latent: "latent", | |
| negative: "negative", | |
| sampler: "sampler", | |
| decode: "decode", | |
| save: "save", | |
| } as const; | |
| /** | |
| * Builds the API-format prompt for ComfyUI's /prompt endpoint. | |
| * | |
| * The wiring mirrors the official "YuE2: Text to Music" and "YuE2: Music Cover" | |
| * templates: CheckpointLoaderSimple feeds YuE2GenerateMusic, whose seconds | |
| * output sizes the empty latent, and KSampler runs the flow model at cfg 1 | |
| * with dpm_2 / sgm_uniform. SheetSage2AudioToABC emits a list, so a PreviewAny | |
| * sits between it and the abc input, exactly as the template does. That relay | |
| * also surfaces the score in /history, which is how the app reads it back. | |
| */ | |
| export function buildPrompt(req: GenerationRequest, seed: number): ComfyPrompt { | |
| const a = req.advanced; | |
| const prompt: ComfyPrompt = {}; | |
| prompt[NODE.checkpoint] = { | |
| class_type: "CheckpointLoaderSimple", | |
| inputs: { ckpt_name: req.checkpoint }, | |
| _meta: { title: "YuE2 checkpoint" }, | |
| }; | |
| // The renderer adapter, when asked for. It changes the acoustic prior the | |
| // flow model samples from, which is why it sits on the MODEL branch alone | |
| // and leaves CLIP on the checkpoint. VAE is untouched either way. | |
| // The two adapters chain. Each loader has to be handed both branches, but | |
| // each advances only the one it actually modifies, so nothing downstream | |
| // reads a branch back out of a loader that was set to zero strength on it. | |
| // The repos are explicit that the other branch is untouched; this keeps the | |
| // graph saying the same thing. VAE never enters the chain. | |
| let modelRef: [string, number] = [NODE.checkpoint, 0]; | |
| let clipRef: [string, number] = [NODE.checkpoint, 1]; | |
| const narStrength = clampStrength(req.realAudioStrength); | |
| const instStrength = clampStrength(req.instrumentalStrength); | |
| if (req.realAudioMatch && narStrength > 0) { | |
| prompt[NODE.loraNar] = { | |
| class_type: "LoraLoader", | |
| inputs: { | |
| model: modelRef, | |
| clip: clipRef, | |
| lora_name: REAL_AUDIO_LORA, | |
| strength_model: narStrength, | |
| strength_clip: 0, | |
| }, | |
| _meta: { title: "Real audio renderer" }, | |
| }; | |
| modelRef = [NODE.loraNar, 0]; | |
| } | |
| if (req.instrumentalArrangement && instStrength > 0 && instrumentalRequirementsMet(req)) { | |
| prompt[NODE.loraInst] = { | |
| class_type: "LoraLoader", | |
| inputs: { | |
| model: modelRef, | |
| clip: clipRef, | |
| lora_name: INSTRUMENTAL_LORA, | |
| strength_model: 0, | |
| strength_clip: instStrength, | |
| }, | |
| _meta: { title: "Instrumental arrangement" }, | |
| }; | |
| clipRef = [NODE.loraInst, 1]; | |
| } | |
| let abcInput: unknown = req.abc; | |
| if (req.scoreSource === "reference" && req.referenceFilename) { | |
| prompt[NODE.audioEncoder] = { | |
| class_type: "AudioEncoderLoader", | |
| inputs: { audio_encoder_name: req.audioEncoder }, | |
| _meta: { title: "SheetSage2" }, | |
| }; | |
| prompt[NODE.loadAudio] = { | |
| class_type: "LoadAudio", | |
| inputs: { audio: req.referenceFilename }, | |
| _meta: { title: "Reference recording" }, | |
| }; | |
| // SheetSage2 quantizes what it hears onto a subbeat grid and gives up on | |
| // notes it cannot place, which a long or busy recording eventually | |
| // contains. Trimming first both avoids that and cuts transcription time. | |
| let audioSource: [string, number] = [NODE.loadAudio, 0]; | |
| if (req.referenceStart > 0 || req.referenceDuration !== null) { | |
| prompt[NODE.trim] = { | |
| class_type: "TrimAudioDuration", | |
| inputs: { | |
| audio: [NODE.loadAudio, 0], | |
| start_index: req.referenceStart, | |
| // The node wants a number; a very large one means "to the end". | |
| duration: req.referenceDuration ?? 100000, | |
| }, | |
| _meta: { title: "Trim reference" }, | |
| }; | |
| audioSource = [NODE.trim, 0]; | |
| } | |
| prompt[NODE.transcribe] = { | |
| class_type: "SheetSage2AudioToABC", | |
| inputs: { | |
| audio_encoder: [NODE.audioEncoder, 0], | |
| audio: audioSource, | |
| mode: req.mode, | |
| }, | |
| _meta: { title: "Transcribe reference" }, | |
| }; | |
| prompt[NODE.abcRelay] = { | |
| class_type: "PreviewAny", | |
| inputs: { source: [NODE.transcribe, 0] }, | |
| _meta: { title: "Transcribed score" }, | |
| }; | |
| abcInput = [NODE.abcRelay, 0]; | |
| } else if (req.scoreSource === "model") { | |
| prompt[NODE.planner] = { | |
| class_type: "YuE2GenerateABC", | |
| inputs: { | |
| clip: clipRef, | |
| style: req.style, | |
| lyrics: req.lyrics, | |
| seed, | |
| mode: req.mode, | |
| max_abc_tokens: a.maxAbcTokens, | |
| temperature: a.abcTemperature, | |
| top_p: a.abcTopP, | |
| top_k: a.abcTopK, | |
| repetition_penalty: a.abcRepetitionPenalty, | |
| penalty_window: 100, | |
| }, | |
| _meta: { title: "Write the score" }, | |
| }; | |
| prompt[NODE.plannerRelay] = { | |
| class_type: "PreviewAny", | |
| inputs: { source: [NODE.planner, 0] }, | |
| _meta: { title: "Generated score" }, | |
| }; | |
| abcInput = [NODE.plannerRelay, 0]; | |
| } | |
| prompt[NODE.music] = { | |
| class_type: "YuE2GenerateMusic", | |
| inputs: { | |
| clip: clipRef, | |
| style: req.style, | |
| lyrics: req.lyrics, | |
| abc: abcInput, | |
| seed, | |
| mode: req.mode, | |
| max_duration: req.maxDuration, | |
| temperature: a.temperature, | |
| top_p: a.topP, | |
| top_k: a.topK, | |
| repetition_penalty: a.repetitionPenalty, | |
| cfg_scale: a.cfgScale, | |
| }, | |
| _meta: { title: "Condition on style and lyrics" }, | |
| }; | |
| prompt[NODE.latent] = { | |
| class_type: "EmptyYuE2LatentAudio", | |
| inputs: { seconds: [NODE.music, 1], batch_size: 1 }, | |
| _meta: { title: "Audio canvas" }, | |
| }; | |
| prompt[NODE.negative] = { | |
| class_type: "ConditioningZeroOut", | |
| inputs: { conditioning: [NODE.music, 0] }, | |
| _meta: { title: "Null conditioning" }, | |
| }; | |
| prompt[NODE.sampler] = { | |
| class_type: "KSampler", | |
| inputs: { | |
| model: modelRef, | |
| positive: [NODE.music, 0], | |
| negative: [NODE.negative, 0], | |
| latent_image: [NODE.latent, 0], | |
| seed, | |
| steps: a.steps, | |
| cfg: 1, | |
| sampler_name: a.samplerName, | |
| scheduler: a.scheduler, | |
| denoise: 1, | |
| }, | |
| _meta: { title: "Render" }, | |
| }; | |
| prompt[NODE.decode] = req.tiledDecode | |
| ? { | |
| class_type: "VAEDecodeAudioTiled", | |
| inputs: { | |
| samples: [NODE.sampler, 0], | |
| vae: [NODE.checkpoint, 2], | |
| tile_size: 1920, | |
| overlap: 128, | |
| }, | |
| _meta: { title: "Decode (tiled)" }, | |
| } | |
| : { | |
| class_type: "VAEDecodeAudio", | |
| inputs: { samples: [NODE.sampler, 0], vae: [NODE.checkpoint, 2] }, | |
| _meta: { title: "Decode" }, | |
| }; | |
| // SaveAudioAdvanced takes quality only for the lossy formats. | |
| const format = req.outputFormat ?? "flac"; | |
| const saveInputs: Record<string, unknown> = { | |
| audio: [NODE.decode, 0], | |
| filename_prefix: "audio/yue2-studio", | |
| format, | |
| }; | |
| if (format === "mp3") saveInputs.quality = "320k"; | |
| if (format === "opus") saveInputs.quality = "192k"; | |
| prompt[NODE.save] = { | |
| class_type: "SaveAudioAdvanced", | |
| inputs: saveInputs, | |
| _meta: { title: "Save" }, | |
| }; | |
| return prompt; | |
| } | |
| function clampStrength(value: number): number { | |
| if (!Number.isFinite(value)) return 0; | |
| return Math.min(1.5, Math.max(0, value)); | |
| } | |
| /** Every node class this app needs ComfyUI to have installed. */ | |
| export const REQUIRED_NODES = [ | |
| "CheckpointLoaderSimple", | |
| "YuE2GenerateABC", | |
| "YuE2GenerateMusic", | |
| "EmptyYuE2LatentAudio", | |
| "SheetSage2AudioToABC", | |
| "AudioEncoderLoader", | |
| "LoadAudio", | |
| "TrimAudioDuration", | |
| "ConditioningZeroOut", | |
| "KSampler", | |
| "VAEDecodeAudio", | |
| "VAEDecodeAudioTiled", | |
| "SaveAudioAdvanced", | |
| "PreviewAny", | |
| "LoraLoader", | |
| ]; | |