Enable YuE2 generation from its model card

packages/civitai-shared/src/basemodel.constants.ts

@@ -1240,8 +1240,7 @@ export const ecosystemSupport: EcosystemSupport[] = [
   // the graph exposes no resources node, so advertising LoRA support would offer
   // resources the form cannot send.
   { ecosystemId: ECO.MiniMaxMusic3, supportType: 'generation', modelTypes: checkpointOnly },
-  // The recipe resolves its own weights; no Civitai resource selection yet.
-  { ecosystemId: ECO.YuE2, supportType: 'generation', modelTypes: [] },
+  { ecosystemId: ECO.YuE2, supportType: 'generation', modelTypes: checkpointOnly },
 
   // PolyGen - remote 3D generator (Meshy via Fal). No Civitai checkpoint/LoRA;
   // entry exists so the unified generator picker can route 3D-Models workflows
@@ -1711,6 +1710,13 @@ export const ecosystemSettings: EcosystemSettings[] = [
       modelLocked: true,
     },
   },
+  {
+    ecosystemId: ECO.YuE2,
+    defaults: {
+      model: { id: 3337846 },
+      modelLocked: true,
+    },
+  },
   {
     ecosystemId: ECO.AceAudio,
     defaults: {

src/components/Model/ModelVersions/ModelVersionDetails.tsx

@@ -314,10 +314,10 @@ function ModelVersionDetailsContent({ model, version, image, onFavoriteClick }:
   const isDraft = version?.status === ModelStatus.Draft;
 
   // const shouldOmit = [1562709, 1672021, 1669468].includes(model.id) && !user?.isModerator;
-  // Drafts hide the action, except for owners/mods on ExternalGeneration versions: those carry no
-  // weights, so generating is the only way to check the wiring before publishing.
+  // Owners/mods can test covered drafts before publishing, including hosted weights.
+  // version.canGenerate below still enforces coverage, ecosystem support, and generation gates.
   const couldGenerate =
-    (!isDraft || (isExternalGeneration && isOwnerOrMod)) &&
+    (!isDraft || isOwnerOrMod) &&
     isSelectableInGenerator &&
     features.imageGeneration &&
     // !shouldOmit &&

src/server/services/orchestrator/__tests__/yue2.handler.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, it } from 'vitest';
 import {
   getEcosystemDisplayItems,
+  isBaseModelGenerationSupported,
   isSelfHostedEcosystem,
 } from '@civitai/shared/basemodel.constants';
 import { generationGraph } from '~/shared/data-graph/generation/generation-graph';
@@ -11,6 +12,8 @@ import { getEcosystemStates as getFormEcosystemStates } from '~/shared/form-grap
 import { createEcosystemStepInput } from '../ecosystems';
 import { createFormGraphStepInput } from '../form-graph';
 import { formatStepOutputs, type GenerationHandlerCtx } from '../orchestration-new.service';
+import { mapDataToGraphInput } from '../legacy-metadata-mapper';
+import type { GenerationResource } from '~/shared/types/generation.types';
 
 const ext: GenerationCtx = {
   limits: { maxQuantity: 4, maxResources: 9, vidQuantity: 1 },
@@ -48,6 +51,31 @@ describe.each([
     dispatch: createFormGraphStepInput,
   },
 ])('YuE2 $name', ({ parse, dispatch }) => {
+  it('opens the official model card in the music generator with v2 selected', () => {
+    const model = {
+      id: 3337846,
+      baseModel: 'YuE2',
+      model: { id: 2944296, type: 'Checkpoint' },
+    };
+    const params = mapDataToGraphInput({}, [model as GenerationResource]);
+    expect(params).toMatchObject({ ecosystem: 'YuE2', workflow: 'txt2music' });
+    const parsed = parse({ ...params, model, prompt: 'A hopeful synth-pop song' });
+    expect(parsed.success).toBe(true);
+    if (!parsed.success) return;
+    expect(parsed.data).toMatchObject({
+      ecosystem: 'YuE2',
+      workflow: 'txt2music',
+      model: { id: 3337846 },
+      yue2MusicMode: 'simple',
+    });
+  });
+
+  it('selects the official v2 checkpoint when starting from the ecosystem picker', () => {
+    const parsed = parse(base);
+    expect(parsed.success).toBe(true);
+    if (parsed.success) expect(parsed.data).toMatchObject({ model: { id: 3337846 } });
+  });
+
   async function submit(overrides: Record<string, unknown> = {}) {
     const parsed = parse({ ...base, ...overrides });
     if (!parsed.success) throw new Error(JSON.stringify(parsed.errors));
@@ -211,4 +239,6 @@ it('offers YuE2 in the audio picker and applies self-hosted availability', () =>
     expect.arrayContaining([expect.objectContaining({ key: 'YuE2', compatible: true })])
   );
   expect(isSelfHostedEcosystem('YuE2')).toBe(true);
+  expect(isBaseModelGenerationSupported('YuE2', 'Checkpoint')).toBe(true);
+  expect(isBaseModelGenerationSupported('YuE2', 'LORA')).toBe(false);
 });

src/shared/data-graph/generation/yue2-graph.ts

@@ -6,7 +6,16 @@ import {
   yue2Steps,
 } from '~/shared/constants/yue2.constants';
 import type { GenerationCtx } from './context';
-import { createTextEditorGraph, enumNode, seedNode, sliderNode, textNode } from './common';
+import {
+  createCheckpointGraph,
+  createTextEditorGraph,
+  enumNode,
+  seedNode,
+  sliderNode,
+  textNode,
+} from './common';
+
+export const yue2VersionIds = { v2: 3337846 } as const;
 
 type YuE2ModeCtx = {
   ecosystem: string;
@@ -40,6 +49,14 @@ const custom = new DataGraph<YuE2ModeCtx, GenerationCtx>()
   ]);
 
 export const yue2Graph = new DataGraph<{ ecosystem: string; workflow: string }, GenerationCtx>()
+  .merge(
+    () =>
+      createCheckpointGraph({
+        versions: { options: [{ label: 'v2', value: yue2VersionIds.v2 }] },
+        defaultModelId: yue2VersionIds.v2,
+      }),
+    []
+  )
   .node('duration', sliderNode({ ...yue2Duration, defaultValue: yue2Duration.default }))
   .node('seed', seedNode())
   .node('yue2MusicMode', enumNode({ options: yue2MusicModeOptions, defaultValue: 'simple' }))

src/shared/form-graph/generation/audio/yue2.graph.ts

@@ -7,8 +7,11 @@ import {
   yue2Steps,
 } from '~/shared/constants/yue2.constants';
 import { SEED, enumDef, sliderDef, textDef } from '../defs';
+import { checkpointDef } from '../checkpoint';
 import { familyScope, type FamilyExt } from '../shared';
 
+export const yue2VersionIds = { v2: 3337846 } as const;
+
 type YuE2Ext = FamilyExt & { yue2MusicMode?: 'simple' | 'custom' };
 
 const simple = defineGraph<YuE2Ext>().field('prompt', {
@@ -34,6 +37,15 @@ const custom = defineGraph<YuE2Ext>()
   .field('yue2Abc', ({ yue2Mode }) => (yue2Mode !== 'off' ? textDef('yue2Abc') : null));
 
 export const yue2 = defineGraph<FamilyExt>({ scope: familyScope })
+  .field('model', ({ _ext }) =>
+    checkpointDef({
+      ecosystem: _ext.ecosystem,
+      workflow: _ext.workflow,
+      ext: _ext,
+      versions: { options: [{ label: 'v2', value: yue2VersionIds.v2 }] },
+      defaultModelId: yue2VersionIds.v2,
+    })
+  )
   .field('duration', sliderDef(yue2Duration))
   .field('seed', SEED)
   .field('yue2MusicMode', enumDef({ options: yue2MusicModeOptions, default: 'simple' }))