Spaces:
Sleeping
Sleeping
| import { useState, useCallback, useRef, useEffect } from 'react' | |
| const DEFAULT_PARAMS = { | |
| musicgen: { | |
| prompt: '', | |
| duration: 15, | |
| seed: -1, | |
| }, | |
| 'ace-step': { | |
| task_type: 'repaint', | |
| caption: '', | |
| lyrics: '[Instrumental]', | |
| bpm: null, | |
| duration: 30, | |
| audio_cover_strength: 0.8, | |
| model_variant: 'xl-turbo', | |
| inference_steps: 15, | |
| seed: -1, | |
| scheduler: 'ode', | |
| cfg_scale: null, | |
| timestep_shift: null, | |
| }, | |
| } | |
| export function useGeneration({ | |
| sessionId, | |
| generate, | |
| extractSeed, | |
| generateAceStep, | |
| checkModels, | |
| }) { | |
| const [selectedModel, setSelectedModel] = useState('musicgen') | |
| const [modelParams, setModelParams] = useState(() => JSON.parse(JSON.stringify(DEFAULT_PARAMS))) | |
| const [isGenerating, setIsGenerating] = useState(false) | |
| const [generationResult, setGenerationResult] = useState(null) | |
| const [generationProgress, setGenerationProgress] = useState(null) | |
| const [modelAvailability, setModelAvailability] = useState(null) | |
| const [seedData, setSeedData] = useState(null) | |
| const abortRef = useRef(null) | |
| // Check model availability and clear stale results when session changes | |
| useEffect(() => { | |
| // Clear previous generation when session changes | |
| setGenerationResult(null) | |
| setSeedData(null) | |
| setGenerationProgress(null) | |
| if (checkModels) { | |
| checkModels().then(data => { | |
| if (data?.models) { | |
| const avail = {} | |
| data.models.forEach(m => { avail[m.name] = m }) | |
| setModelAvailability(avail) | |
| } | |
| }) | |
| } | |
| }, [sessionId, checkModels]) | |
| const setModelParam = useCallback((model, key, value) => { | |
| setModelParams(prev => ({ | |
| ...prev, | |
| [model]: { ...prev[model], [key]: value }, | |
| })) | |
| }, []) | |
| const resetParams = useCallback((model) => { | |
| setModelParams(prev => ({ | |
| ...prev, | |
| [model]: JSON.parse(JSON.stringify(DEFAULT_PARAMS[model])), | |
| })) | |
| }, []) | |
| // Extract seed from region | |
| const doExtractSeed = useCallback(async (regionStart, regionEnd) => { | |
| if (!extractSeed) return null | |
| const result = await extractSeed(regionStart, regionEnd) | |
| if (result) { | |
| setSeedData(result) | |
| // Pre-fill BPM from analysis | |
| if (result.analysis?.bpm) { | |
| setModelParams(prev => ({ | |
| ...prev, | |
| 'ace-step': { ...prev['ace-step'], bpm: result.analysis.bpm }, | |
| })) | |
| } | |
| } | |
| return result | |
| }, [extractSeed]) | |
| // Start generation | |
| const startGeneration = useCallback(async (regionStart, regionEnd) => { | |
| setIsGenerating(true) | |
| setGenerationResult(null) | |
| setGenerationProgress({ status: 'starting' }) | |
| try { | |
| let result = null | |
| if (selectedModel === 'musicgen') { | |
| const params = modelParams.musicgen | |
| result = await generate(regionStart, regionEnd, params.duration, params.prompt || null) | |
| } else if (selectedModel === 'ace-step') { | |
| if (!seedData?.seed_id) { | |
| // Auto-extract seed if not done | |
| const extracted = await doExtractSeed(regionStart, regionEnd) | |
| if (!extracted) throw new Error('Failed to extract seed') | |
| } | |
| const params = { ...modelParams['ace-step'] } | |
| params.seed_id = seedData.seed_id | |
| result = await generateAceStep(params) | |
| } | |
| if (result) { | |
| setGenerationResult({ ...result, model: selectedModel }) | |
| setGenerationProgress({ status: 'done' }) | |
| } else { | |
| setGenerationProgress({ status: 'error' }) | |
| } | |
| return result | |
| } catch (err) { | |
| setGenerationProgress({ status: 'error', message: err.message }) | |
| return null | |
| } finally { | |
| setIsGenerating(false) | |
| } | |
| }, [selectedModel, modelParams, seedData, generate, generateAceStep, doExtractSeed]) | |
| const cancelGeneration = useCallback(() => { | |
| if (abortRef.current) { | |
| abortRef.current.abort() | |
| abortRef.current = null | |
| } | |
| setIsGenerating(false) | |
| setGenerationProgress(null) | |
| }, []) | |
| return { | |
| selectedModel, | |
| setSelectedModel, | |
| modelParams, | |
| setModelParam, | |
| resetParams, | |
| isGenerating, | |
| generationResult, | |
| setGenerationResult, | |
| generationProgress, | |
| modelAvailability, | |
| seedData, | |
| setSeedData, | |
| doExtractSeed, | |
| startGeneration, | |
| cancelGeneration, | |
| } | |
| } | |