diff --git a/packages/civitai-shared/src/basemodel.constants.ts b/packages/civitai-shared/src/basemodel.constants.ts index 80947c26b8..f01f45bd02 100644 --- a/packages/civitai-shared/src/basemodel.constants.ts +++ b/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: { diff --git a/src/components/Model/ModelVersions/ModelVersionDetails.tsx b/src/components/Model/ModelVersions/ModelVersionDetails.tsx index a017f1393b..54128c5aa1 100644 --- a/src/components/Model/ModelVersions/ModelVersionDetails.tsx +++ b/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 && diff --git a/src/server/services/orchestrator/__tests__/yue2.handler.test.ts b/src/server/services/orchestrator/__tests__/yue2.handler.test.ts index ba6fc0e576..7a3d830704 100644 --- a/src/server/services/orchestrator/__tests__/yue2.handler.test.ts +++ b/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 = {}) { 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); }); diff --git a/src/shared/data-graph/generation/yue2-graph.ts b/src/shared/data-graph/generation/yue2-graph.ts index b7eafc9abe..540097bcd5 100644 --- a/src/shared/data-graph/generation/yue2-graph.ts +++ b/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() ]); 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' })) diff --git a/src/shared/form-graph/generation/audio/yue2.graph.ts b/src/shared/form-graph/generation/audio/yue2.graph.ts index 16f16e3b5e..69792a3023 100644 --- a/src/shared/form-graph/generation/audio/yue2.graph.ts +++ b/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().field('prompt', { @@ -34,6 +37,15 @@ const custom = defineGraph() .field('yue2Abc', ({ yue2Mode }) => (yue2Mode !== 'off' ? textDef('yue2Abc') : null)); export const yue2 = defineGraph({ 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' }))