Initial experimental support for yue2

2026-09-14 05:07:11 -06:00
1 parent 82137c8ca0
commit e43f138ae4
@@ -13,7 +13,7 @@ import { startQueue } from '@/utils/queue';
 import useGPUInfo from '@/hooks/useGPUInfo';
 import usePollLoop from '@/hooks/usePollLoop';
 import { defaultInferenceJobConfig } from '@/helpers/inferenceJobConfig';
-import { encodeFilePathForUrl } from '@/utils/basic';
+import { encodeFilePathForUrl, objToTags, tagsToObj } from '@/utils/basic';
 import { latentToImage, payloadToFloat32, readEngineFrames, PreviewInfo } from '@/utils/engineStream';
 import { isMac } from '@/helpers/basic';
 import { modelArchs, getGenerateDefaults, GenerateDefaults } from '@/app/jobs/new/options';
@@ -842,7 +842,19 @@ function GeneratePageInner() {
                          className={`h-16 w-16 rounded-md overflow-hidden border-2 ${isSel ? 'border-blue-500' : 'border-transparent hover:border-gray-600'} bg-gray-800 block`}
                        >
                          {r.kind === 'audio' ? (
-                            <div className="w-full h-full flex items-center justify-center text-gray-400 text-[10px]">audio</div>
+                            <div className="relative w-full h-full flex items-center justify-center text-gray-400 text-[10px]">
+                              audio
+                              {/* embedded waveform cover; falls back to the label when the file has none */}
+                              {/* eslint-disable-next-line @next/next/no-img-element */}
+                              <img
+                                src={`/api/audio/art/${encodeFilePathForUrl(r.path)}`}
+                                alt=""
+                                className="absolute inset-0 w-full h-full object-cover"
+                                onError={e => {
+                                  (e.target as HTMLImageElement).style.display = 'none';
+                                }}
+                              />
+                            </div>
                          ) : r.kind === 'video' ? (
                            <video src={resultUrl(r)} muted className="w-full h-full object-cover" />
                          ) : (
@@ -1025,16 +1037,35 @@ function GeneratePageInner() {
                  </div>
                </Card>
                <Card title="Prompt" open={cardOpen('prompt')} onToggle={() => toggleCard('prompt')}>
-                  <TextAreaInput label="Prompt" value={sample.prompt || ''} onChange={v => setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} />
-                  {(sample.guidance_scale ?? 4) > 1 && (
+                  {entry?.sampleTags ? (
+                    // tagged prompt archs (audio): one field per tag, stored as the tagged prompt string
+                    Object.entries(entry.sampleTags).map(([tagKey, tag]) => {
+                      const tags = tagsToObj(sample.prompt || '');
+                      const setTag = (v: any) => setSample(s => ({ ...s, prompt: objToTags({ ...tagsToObj(s.prompt || ''), [tagKey]: v }) }));
+                      const value = tags[tagKey] ?? '';
+                      if (tag.type === 'multiline') return <TextAreaInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                      if (tag.type === 'number') return <NumberInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                      return <TextInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                    })
+                  ) : (
+                    <TextAreaInput label="Prompt" value={sample.prompt || ''} onChange={v => setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} />
+                  )}
+                  {(sample.guidance_scale ?? 4) > 1 && modality !== 'audio' && (
                    <TextAreaInput label="Negative prompt" value={sample.negative_prompt || ''} onChange={v => setSample(s => ({ ...s, negative_prompt: v }))} />
                  )}
                  <div className="grid grid-cols-3 gap-2">
-                    <NumberInput label="Width" value={sample.width ?? 1024} onChange={v => setSample(s => ({ ...s, width: v }))} min={64} max={4096} />
-                    <NumberInput label="Height" value={sample.height ?? 1024} onChange={v => setSample(s => ({ ...s, height: v }))} min={64} max={4096} />
+                    {modality !== 'audio' && (
+                      <>
+                        <NumberInput label="Width" value={sample.width ?? 1024} onChange={v => setSample(s => ({ ...s, width: v }))} min={64} max={4096} />
+                        <NumberInput label="Height" value={sample.height ?? 1024} onChange={v => setSample(s => ({ ...s, height: v }))} min={64} max={4096} />
+                      </>
+                    )}
                    <NumberInput label="Steps" value={sample.num_inference_steps ?? 25} onChange={v => setSample(s => ({ ...s, num_inference_steps: v }))} min={1} max={200} />
                    <NumberInput label="Guidance" value={sample.guidance_scale ?? 4} onChange={v => setSample(s => ({ ...s, guidance_scale: v }))} min={1} max={30} />
                    <NumberInput label="Seed (-1 random)" value={sample.seed ?? -1} onChange={v => setSample(s => ({ ...s, seed: v }))} min={-1} max={4294967295} />
+                    {modality === 'audio' && sample.duration !== undefined && (
+                      <NumberInput label="Duration (s)" value={sample.duration} onChange={v => setSample(s => ({ ...s, duration: v }))} min={1} max={600} />
+                    )}
                    {modality === 'video' && (
                      <>
                        <NumberInput label="Frames" value={sample.num_frames ?? 33} onChange={v => setSample(s => ({ ...s, num_frames: v }))} min={1} max={1000} />
@@ -1524,6 +1524,18 @@ export default function SimpleJob({
                </div>
              )}

+              {modelArch?.additionalSections?.includes('sample.duration') && (
+                <div>
+                  <NumberInput
+                    label="Duration (seconds)"
+                    value={jobConfig.config.process[0].sample.duration ?? 120}
+                    onChange={value => setJobConfig(value, 'config.process[0].sample.duration')}
+                    placeholder="eg. 120"
+                    min={1}
+                    required
+                  />
+                </div>
+              )}
              <div>
                <NumberInput
                  label="Seed"
@@ -2,7 +2,7 @@ import React from 'react';
 import Link from 'next/link';
 import { GroupedSelectOption, SelectOption, JobConfig, ConfigDoc } from '@/types';
 import { defaultSliderConfig } from './jobConfig';
-import { defaultAudioSampleConfig, defaultSampleConfig, defaultIdeogramSamplesConfig } from '@/helpers/defaultSamples';
+import { defaultAudioSampleConfig, defaultSampleConfig, defaultIdeogramSamplesConfig, defaultYue2SampleConfig } from '@/helpers/defaultSamples';

 type Control = 'depth' | 'line' | 'pose' | 'inpaint';

@@ -27,6 +27,7 @@ type AdditionalSections =
  | 'datasets.auto_frame_count'
  | 'sample.ctrl_img'
  | 'sample.multi_ctrl_imgs'
+  | 'sample.duration'
  | 'train.audio_loss_multiplier'
  | 'datasets.num_frames'
  | 'model.multistage'
@@ -80,6 +81,8 @@ export interface GenerateDefaults {
  sample: { [key: string]: any };
  needsControlImage: boolean;
  sizeLocked: boolean;
+  /** structured prompt fields (audio models): the prompt is their tagged form */
+  sampleTags?: SampleTags;
 }

 export interface ModelArch {
@@ -1370,6 +1373,31 @@ export const modelArchs: ModelArch[] = [
    disableSections: ['network.conv'],
    additionalSections: ['sample.multi_ctrl_imgs', 'model.low_vram', 'model.layer_offloading'],
  },
+  {
+    name: 'yue2',
+    label: 'YuE2',
+    group: 'audio',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': [
+        'Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors',
+        defaultNameOrPath,
+      ],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [false, false],
+      'config.process[0].model.low_vram': [false, false],
+      'config.process[0].train.unload_text_encoder': [false, false],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['sigmoid', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+      'config.process[0].sample': [defaultYue2SampleConfig, defaultSampleConfig],
+      'config.process[0].datasets[x].cache_latents_to_disk': [true, true],
+    },
+    // native YuE2 prompt: style text, a [Lyrics] line, the lyrics
+    hasMultiLinePrompts: true,
+    disableSections: ['network.conv'],
+    additionalSections: ['model.low_vram', 'sample.duration'],
+  },
  {
    name: 'ace_step_15',
    label: 'ACE-Step 1.5',
@@ -1832,13 +1860,14 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => {
      if (sc.guidance_scale !== undefined) sample.guidance_scale = sc.guidance_scale;
      if (sc.num_frames) sample.num_frames = sc.num_frames;
      if (sc.fps) sample.fps = sc.fps;
+      if (sc.duration) sample.duration = sc.duration;
      if (sc.neg) sample.negative_prompt = sc.neg;
    } else if (key.startsWith(SAMPLE_PREFIX + '.')) {
      const field = key.slice(SAMPLE_PREFIX.length + 1);
      if (value === undefined || value === null || value === '') continue;
      if (field === 'sample_steps') sample.num_inference_steps = value;
      else if (field === 'neg') sample.negative_prompt = value;
-      else if (['width', 'height', 'guidance_scale', 'num_frames', 'fps'].includes(field)) sample[field] = value;
+      else if (['width', 'height', 'guidance_scale', 'num_frames', 'fps', 'duration'].includes(field)) sample[field] = value;
    }
  }
  // the engine defaults to convrot8; the training default qtype is the
@@ -1861,5 +1890,6 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => {
    sample: { ...sample, ...(gen.sample || {}) },
    needsControlImage: gen.needsControlImage ?? (sections.includes('sample.ctrl_img') || sections.includes('sample.multi_ctrl_imgs')),
    sizeLocked: gen.sizeLocked ?? false,
+    sampleTags: arch.sampleTags,
  };
 };
@@ -977,3 +977,22 @@ export const defaultIdeogramSamplesConfig: SampleConfig = {
  num_frames: 1,
  fps: 1,
 }
+
+// YuE2 takes its native prompt layout (style text, a [Lyrics] line, the lyrics):
+// derive it from the shared audio samples instead of duplicating the lyrics
+const tagValue = (prompt: string, tag: string) => {
+  const m = prompt.match(new RegExp(`<${tag}>([\\s\\S]*?)</${tag}>`));
+  return m ? m[1].trim() : '';
+};
+
+export const defaultYue2SampleConfig: SampleConfig = {
+  ...defaultAudioSampleConfig,
+  samples: defaultAudioSampleConfig.samples.map(s => ({
+    ...s,
+    prompt: `${tagValue(s.prompt, 'CAPTION')}\n[Lyrics]\n${tagValue(s.prompt, 'LYRICS')}\n`,
+  })),
+  guidance_scale: 1,
+  sample_steps: 32,
+  // max seconds per sample; the AR stops earlier when the song ends
+  duration: 120,
+};
@@ -229,6 +229,7 @@ export interface SampleItem {
  sample_steps?: number;
  fps?: number;
  num_frames?: number;
+  duration?: number;
  ctrl_img?: string | null;
  ctrl_idx?: number;
  network_multiplier?: number;
@@ -252,6 +253,7 @@ export interface SampleConfig {
  sample_steps: number;
  num_frames: number;
  fps: number;
+  duration?: number;
 }

 export interface LoggingConfig {