import { useRouter } from "next/router"; import { unmute } from "../external/unmute"; import { useCallback, useEffect, useState } from "react"; import * as Tone from "tone"; import AudioPlayer from "../components/AudioPlayer"; import FallingBehindWarning from "../components/FallingBehindWarning"; import PageHead from "../components/PageHead"; import Share from "../components/Share"; import Settings from "../components/Settings"; import ModelInference from "../components/ModelInference"; import Pause from "../components/Pause"; import PromptPanel from "../components/PromptPanel"; import ThreeCanvas from "../components/ThreeCanvas"; import { samplePrompts, initialSeeds, initialSeedImageMap } from "../samplePrompts"; import { AppState, InferenceInput, InferenceResult, PromptInput, } from "../types"; function getRandomFromArray(arr: any[], n: number) { var result = new Array(n), len = arr.length, taken = new Array(len); if (n > len) throw new RangeError("getRandom: more elements taken than available"); while (n--) { var x = Math.floor(Math.random() * len); result[n] = arr[x in taken ? taken[x] : x]; taken[x] = --len in taken ? taken[len] : len; } return result; } export default function Home() { // High-level state of the prompts const [appState, setAppState] = useState(AppState.UNINITIALIZED); // Whether playback is paused const [paused, setPaused] = useState(true); // Current interpolation parameters const [alpha, setAlpha] = useState(0.0); const [alphaRollover, setAlphaRollover] = useState(false); const [alphaVelocity, setAlphaVelocity] = useState(0.25); // Settings const [denoising, setDenoising] = useState(0.75); const [seedImageId, setSeedImageId] = useState( initialSeeds[Math.floor(Math.random() * initialSeeds.length)] ); const [seed, setSeed] = useState( initialSeedImageMap[seedImageId][ Math.floor(Math.random() * initialSeedImageMap[seedImageId].length) ] ); // Prompts shown on screen and maintained by the prompt panel const [promptInputs, setPromptInputs] = useState([]); // Model execution results const [inferenceResults, setInferenceResults] = useState( [] ); // Currently playing result, from the audio player const [nowPlayingResult, setNowPlayingResult] = useState(null); // Set the initial seed from the URL if available const router = useRouter(); useEffect(() => { // Make sure params are ready if (!router.isReady) { return; } if (router.query.alphaVelocity) { setAlphaVelocity(parseFloat(router.query.alphaVelocity as string)); } if (router.query.seed) { setSeed(parseInt(router.query.seed as string)); } // Set prompt inputs here so that they are empty and incorrect information isn't shown // until URL params have a chance to be read. let initPromptInputs = [ { prompt: "" }, { prompt: "" }, { prompt: "" }, { prompt: "" }, { prompt: "" }, { prompt: "" }, ]; // Set random initial prompts const samples = getRandomFromArray(samplePrompts, 4); for (let i = 0; i < 4; i++) { initPromptInputs[i].prompt = samples[i]; } if (router.query.prompt) { initPromptInputs[3].prompt = router.query.prompt as string; } setPromptInputs(initPromptInputs); if (router.query.denoising) { setDenoising(parseFloat(router.query.denoising as string)); } if (router.query.seedImageId) { setSeedImageId(router.query.seedImageId as string); } }, [router.isReady, router.query]); // Set the app state based on the prompt inputs array useEffect(() => { if (!alphaRollover) { return; } setAlphaRollover(false); const upNextPrompt = promptInputs[promptInputs.length - 1].prompt; const endPrompt = promptInputs[promptInputs.length - 2].prompt; let newAppState = appState; if (appState == AppState.SAME_PROMPT) { if (upNextPrompt) { newAppState = AppState.TRANSITION; // swap the last two prompts to bring the upNextPrompt into the next inference call setPromptInputs((prevPromptInputs) => { const newPromptInputs = [...prevPromptInputs]; newPromptInputs[newPromptInputs.length - 1] = { prompt: endPrompt, }; newPromptInputs[newPromptInputs.length - 2] = { prompt: upNextPrompt, }; return newPromptInputs; }); } setSeed(seed + 1); } else if (appState == AppState.TRANSITION) { setPromptInputs([...promptInputs, { prompt: "" }]); if (upNextPrompt) { newAppState = AppState.TRANSITION; } else { newAppState = AppState.SAME_PROMPT; } } if (newAppState != appState) { setAppState(newAppState); } }, [promptInputs, alpha, alphaRollover, appState, seed]); // On any app state change, reset alpha useEffect(() => { console.log("App State: ", appState); if (appState == AppState.UNINITIALIZED) { setAppState(AppState.SAME_PROMPT); } else { setAlpha(0.25); } }, [appState]); // What to do when a new inference result is available const newResultCallback = useCallback( (input: InferenceInput, result: InferenceResult) => { setInferenceResults((prevResults: InferenceResult[]) => { const maxResultCounter = Math.max(...prevResults.map((r) => r.counter)); const lastResult = prevResults.find( (r) => r.counter == maxResultCounter ); const newCounter = lastResult ? lastResult.counter + 1 : 0; result.counter = newCounter; result.input = input; result.played = false; let newAlpha = alpha + alphaVelocity; if (newAlpha > 1 + 1e-3) { newAlpha = newAlpha - 1; setAlphaRollover(true); } setAlpha(newAlpha); return [...prevResults, result]; }); }, [alpha, alphaVelocity] ); // State to handle the timeout for the player to not hog GPU forever. If you are // in SAME_PROMPT for this long, it will pause the player and bring up an alert. const timeoutIncrement = 600.0; const [timeoutPlayerTime, setTimeoutPlayerTime] = useState(timeoutIncrement); const nowPlayingCallback = useCallback( (result: InferenceResult, playerTime: number) => { console.log( "Now playing result ", result.counter, ", player time is ", playerTime ); setNowPlayingResult(result); // find the first promptInput that matches the result.input.end.prompt and set it's transitionCounter to the result.counter if not already set setPromptInputs((prevPromptInputs) => { const newPromptInputs = [...prevPromptInputs]; const promptInputIndex = newPromptInputs.findIndex( (p) => p.prompt == result.input.end.prompt ); if (promptInputIndex >= 0) { if (newPromptInputs[promptInputIndex].transitionCounter == null) { newPromptInputs[promptInputIndex].transitionCounter = result.counter; } } return newPromptInputs; }); // set played state for the result to true setInferenceResults((prevResults: InferenceResult[]) => { return prevResults.map((r) => { if (r.counter == result.counter) { r.played = true; } return r; }); }); // Extend the timeout if we're transitioning if (appState == AppState.TRANSITION) { setTimeoutPlayerTime(playerTime + timeoutIncrement); } // If we've hit the timeout, pause and increment the timeout if (playerTime > timeoutPlayerTime) { setTimeoutPlayerTime(playerTime + timeoutIncrement); setPaused(true); setTimeout(() => { if (confirm("Are you still riffing?")) { setPaused(false); } }, 100); } }, [timeoutPlayerTime, appState] ); // Track from the audio player whether we're behind on having new inference results, // in order to display an alert. const [playerIsBehind, setPlayerIsBehind] = useState(false); const playerIsBehindCallback = (isBehind: boolean) => { setPlayerIsBehind(isBehind); }; // TODO(hayk): This is not done yet but it's supposed to clear out future state and prompts // and allow the user to immediately start a new prompt. const resetCallback = () => { console.log("RESET"); setPromptInputs([ promptInputs[0], promptInputs[1], promptInputs[2], { prompt: "" }, { prompt: "" }, { prompt: "" }, ]); if (nowPlayingResult == null) { setInferenceResults([]); } // const counter = nowPlayingResult ? nowPlayingResult.counter : -1; // setInferenceResults(inferenceResults.filter((r) => r.counter <= counter)); // setNowPlayingResult(null); }; useEffect(() => { unmute(Tone.context.rawContext, true, false); }, []); return ( <>
{/* Note, top bg section is used to maintain color in background on ios notch devices */}
window.open("/about", "_blank")} > [RIFFUSION]
{playerIsBehind ? : null}
Tone.Transport.seconds} inferenceResults={inferenceResults} />
{ const newPromptInputs = [...promptInputs]; newPromptInputs[index].prompt = prompt; setPromptInputs(newPromptInputs); }} setPaused={setPaused} resetCallback={resetCallback} />
); }