Spaces:
Sleeping
Sleeping
| const TOTAL_STEPS = 1000; | |
| const modelInput = document.querySelector('#model-input'); | |
| const DEFAULT_CHECKPOINTS = [0, 50, 100, 200, 350, 500, 700, 1000]; | |
| const SIZE = 192; | |
| const PIXEL_COUNT = SIZE * SIZE * 3; | |
| const elements = { | |
| timestepsInput: document.querySelector('#timesteps-input'), applyTimesteps: document.querySelector('#apply-timesteps'), | |
| seedInput: document.querySelector('#seed-input'), randomizeSeed: document.querySelector('#randomize-seed'), configError: document.querySelector('#config-error'), diffusionTrack: document.querySelector('#diffusion-track'), | |
| denoisingTrack: document.querySelector('#denoising-track'), imageUpload: document.querySelector('#image-upload'), modelStatus: document.querySelector('#model-status'), | |
| forwardProgressLabel: document.querySelector('#forward-progress-label'), reverseProgressLabel: document.querySelector('#reverse-progress-label'), | |
| }; | |
| let checkpoints = [...DEFAULT_CHECKPOINTS]; | |
| let sourcePixels = new Float32Array(PIXEL_COUNT); | |
| let forwardStates = []; | |
| let reverseStates = new Map(); | |
| let imageLabel = 'Still life 01'; | |
| let forwardRevealCount = 1; | |
| let forwardCompleted = false; | |
| let reverseNoiseState = null; | |
| let reverseInputMode = 'seed'; | |
| let animationTimer = null; | |
| let generationController = null; | |
| let generationBusy = false; | |
| const progressStartedAt = { forward: null, reverse: null }; | |
| const schedule = makeSchedule(); | |
| function makeSchedule() { | |
| const alphaBars = new Float64Array(TOTAL_STEPS); let product = 1; | |
| for (let step = 0; step < TOTAL_STEPS; step += 1) { const beta = .0001 + (.02 - .0001) * step / (TOTAL_STEPS - 1); product *= 1 - beta; alphaBars[step] = product; } | |
| return alphaBars; | |
| } | |
| function randomNormal(seed) { | |
| let state = seed >>> 0; | |
| return () => { state = (state + 0x6D2B79F5) | 0; let value = state; value = Math.imul(value ^ value >>> 15, value | 1); value ^= value + Math.imul(value ^ value >>> 7, value | 61); const uniform = ((value ^ value >>> 14) >>> 0) / 4294967296; return Math.sqrt(-2 * Math.log(Math.max(uniform, 1e-7))) * Math.cos(2 * Math.PI * uniform); }; | |
| } | |
| function makeSourceArt() { | |
| const canvas = document.createElement('canvas'); canvas.width = SIZE; canvas.height = SIZE; const context = canvas.getContext('2d'); | |
| const sky = context.createLinearGradient(0, 0, SIZE, SIZE); sky.addColorStop(0, '#dce7db'); sky.addColorStop(.58, '#b8c7ad'); sky.addColorStop(1, '#a2b5a0'); context.fillStyle = sky; context.fillRect(0, 0, SIZE, SIZE); | |
| context.fillStyle = '#e9d6b5'; context.beginPath(); context.arc(147, 44, 24, 0, Math.PI * 2); context.fill(); | |
| context.fillStyle = '#52765f'; context.beginPath(); context.moveTo(0, 137); context.quadraticCurveTo(50, 91, 99, 137); context.quadraticCurveTo(145, 109, 192, 132); context.lineTo(192, 192); context.lineTo(0, 192); context.fill(); | |
| context.fillStyle = '#315b48'; context.beginPath(); context.moveTo(0, 159); context.quadraticCurveTo(61, 127, 112, 164); context.quadraticCurveTo(155, 143, 192, 157); context.lineTo(192, 192); context.lineTo(0, 192); context.fill(); | |
| context.fillStyle = '#d8795d'; context.beginPath(); context.ellipse(92, 144, 24, 31, -.22, 0, Math.PI * 2); context.fill(); context.fillStyle = '#e8a37c'; context.beginPath(); context.ellipse(84, 136, 14, 21, -.22, 0, Math.PI * 2); context.fill(); | |
| context.fillStyle = '#496c52'; context.beginPath(); context.ellipse(103, 111, 21, 10, -.62, 0, Math.PI * 2); context.fill(); context.beginPath(); context.ellipse(77, 109, 16, 8, .52, 0, Math.PI * 2); context.fill(); return canvas; | |
| } | |
| function canvasPixels(canvas) { | |
| const rgba = canvas.getContext('2d', { willReadFrequently: true }).getImageData(0, 0, SIZE, SIZE).data; const pixels = new Float32Array(PIXEL_COUNT); | |
| for (let pixel = 0; pixel < SIZE * SIZE; pixel += 1) for (let channel = 0; channel < 3; channel += 1) pixels[pixel * 3 + channel] = rgba[pixel * 4 + channel] / 127.5 - 1; | |
| return pixels; | |
| } | |
| function pixelsToCanvas(pixels, canvas) { | |
| const context = canvas.getContext('2d'); const image = context.createImageData(SIZE, SIZE); | |
| for (let index = 0; index < PIXEL_COUNT; index += 1) image.data[Math.floor(index / 3) * 4 + index % 3] = Math.round((Math.max(-1, Math.min(1, pixels[index])) + 1) * 127.5); | |
| for (let index = 3; index < image.data.length; index += 4) image.data[index] = 255; context.putImageData(image, 0, 0); | |
| } | |
| function drawState(canvas, state) { const context = canvas.getContext('2d'); context.clearRect(0, 0, SIZE, SIZE); if (state instanceof HTMLCanvasElement) context.drawImage(state, 0, 0, SIZE, SIZE); else if (state instanceof Float32Array) pixelsToCanvas(state, canvas); } | |
| function computeForwardStates() { | |
| const random = randomNormal(readSeed()); const noise = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) noise[index] = random(); | |
| forwardStates = checkpoints.map((step) => { if (step === 0) return new Float32Array(sourcePixels); const alphaBar = schedule[step - 1]; const state = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) state[index] = Math.sqrt(alphaBar) * sourcePixels[index] + Math.sqrt(1 - alphaBar) * noise[index]; return state; }); | |
| } | |
| function makeNoiseState(seed) { const random = randomNormal(seed); const state = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) state[index] = random(); return state; } | |
| function stateToPng(state) { const canvas = document.createElement('canvas'); const size = modelInput.value === 'google/ddpm-cifar10-32' ? 32 : 256; canvas.width = size; canvas.height = size; const small = document.createElement('canvas'); small.width = SIZE; small.height = SIZE; drawState(small, state); canvas.getContext('2d').drawImage(small, 0, 0, size, size); return canvas.toDataURL('image/png'); } | |
| function setModelStatus(message) { if (elements.modelStatus) elements.modelStatus.textContent = message; } | |
| function resampleNoise() { | |
| const values = new Uint32Array(1); crypto.getRandomValues(values); const seed = values[0] & 0x7fffffff; elements.seedInput.value = String(seed); reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(seed); reverseStates = new Map(); resetProgress('reverse'); setModelStatus(`Resampled Gaussian noise · seed ${seed}.`); renderTracks(); | |
| } | |
| function useForwardNoise() { | |
| if (!forwardCompleted) return; reverseInputMode = 'forward'; reverseNoiseState = new Float32Array(forwardStates.at(-1)); reverseStates = new Map(); resetProgress('reverse'); setModelStatus('Using the final forward-process noise as the reverse input.'); renderTracks(); | |
| } | |
| function makeActionButton(label, id, action) { const button = document.createElement('button'); button.type = 'button'; button.id = id; button.className = 'process-button'; button.textContent = label; button.addEventListener('click', action); return button; } | |
| function makeCard(step, state, type, index) { | |
| const card = document.createElement('article'); card.className = 'stage-card'; card.dataset.step = String(step); const actionSlot = document.createElement('div'); actionSlot.className = 'action-slot'; | |
| if (type === 'forward' && index >= forwardRevealCount) card.classList.add('is-pending'); | |
| if (type === 'forward' && index === 0) actionSlot.append(makeActionButton('Diffuse', 'diffuse-button', runDiffusion)); | |
| if (type === 'forward' && index === checkpoints.length - 1) { actionSlot.classList.add('status-right'); actionSlot.append(elements.forwardProgressLabel); } | |
| if (type === 'reverse' && index === 0) actionSlot.append(elements.reverseProgressLabel); | |
| if (type === 'reverse' && index === checkpoints.length - 1) { | |
| actionSlot.classList.add('reverse-actions'); | |
| actionSlot.append(makeActionButton('Denoise', 'denoise-button', generateReverse)); | |
| } | |
| const canvasWrap = document.createElement('div'); canvasWrap.className = 'stage-canvas'; const canvas = document.createElement('canvas'); canvas.width = SIZE; canvas.height = SIZE; drawState(canvas, state); canvasWrap.append(canvas); | |
| if (type === 'forward' && index === 0) { const change = document.createElement('button'); change.className = 'change-image'; change.type = 'button'; change.textContent = 'Change image'; change.addEventListener('click', () => elements.imageUpload.click()); canvasWrap.append(change); } | |
| if (type === 'reverse' && index === checkpoints.length - 1) { const noiseActions = document.createElement('div'); noiseActions.className = 'noise-actions'; noiseActions.append(makeActionButton('Resample', 'resample-button', resampleNoise), makeActionButton('Forward noise', 'forward-noise-button', useForwardNoise)); canvasWrap.append(noiseActions); } | |
| const caption = document.createElement('div'); caption.className = 'stage-caption'; | |
| const titleText = type === 'forward' && index === 0 ? imageLabel : step === 0 ? 'Generated image' : step === TOTAL_STEPS ? 'Gaussian noise' : ''; | |
| if (titleText) { const title = document.createElement('strong'); title.textContent = titleText; caption.append(title); } | |
| const time = document.createElement('span'); time.textContent = `t = ${step}`; caption.append(time); | |
| const preview = document.createElement('div'); preview.className = 'hover-preview'; | |
| const previewCanvas = document.createElement('canvas'); | |
| const nativeSize = modelInput.value === 'google/ddpm-cifar10-32' ? 32 : 256; | |
| previewCanvas.width = state instanceof HTMLCanvasElement ? state.width : nativeSize; | |
| previewCanvas.height = state instanceof HTMLCanvasElement ? state.height : nativeSize; | |
| preview.style.width = `${previewCanvas.width + 4}px`; preview.style.height = `${previewCanvas.height + 4}px`; | |
| previewCanvas.style.width = `${previewCanvas.width}px`; previewCanvas.style.height = `${previewCanvas.height}px`; | |
| if (state) { | |
| const context = previewCanvas.getContext('2d'); | |
| if (state instanceof HTMLCanvasElement) context.drawImage(state, 0, 0); | |
| else { const temporary = document.createElement('canvas'); temporary.width = SIZE; temporary.height = SIZE; drawState(temporary, state); context.drawImage(temporary, 0, 0, previewCanvas.width, previewCanvas.height); } | |
| } | |
| preview.append(previewCanvas); | |
| canvas.addEventListener('click', () => { | |
| document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible')); | |
| const rect = canvasWrap.getBoundingClientRect(); | |
| preview.style.left = `${Math.max(8, Math.min(window.innerWidth - previewCanvas.width - 12, rect.left + rect.width / 2 - previewCanvas.width / 2))}px`; | |
| preview.style.top = `${Math.max(8, Math.min(window.innerHeight - previewCanvas.height - 12, rect.top + rect.height / 2 - previewCanvas.height / 2))}px`; | |
| preview.classList.add('is-visible'); | |
| }); | |
| card.append(actionSlot, canvasWrap, caption, preview); return card; | |
| } | |
| function renderTracks() { | |
| elements.diffusionTrack.style.setProperty('--stage-count', checkpoints.length); | |
| elements.denoisingTrack.style.setProperty('--stage-count', checkpoints.length); | |
| elements.diffusionTrack.replaceChildren(...checkpoints.map((step, index) => makeCard(step, forwardStates[index], 'forward', index))); | |
| elements.denoisingTrack.replaceChildren(...checkpoints.map((step, index) => makeCard(step, reverseStates.get(step) || (step === TOTAL_STEPS ? reverseNoiseState : null), 'reverse', index))); | |
| modelInput.disabled = generationBusy; | |
| const button = document.querySelector('#denoise-button'); if (button) { button.disabled = generationBusy; button.textContent = generationBusy ? 'Denoising…' : 'Denoise'; } | |
| const forwardNoiseButton = document.querySelector('#forward-noise-button'); if (forwardNoiseButton) { forwardNoiseButton.disabled = generationBusy || !forwardCompleted; forwardNoiseButton.title = forwardCompleted ? 'Use the last forward-process image' : 'Run Diffuse first'; } | |
| const resampleButton = document.querySelector('#resample-button'); if (resampleButton) resampleButton.disabled = generationBusy; | |
| } | |
| function formatDuration(seconds) { | |
| if (!Number.isFinite(seconds)) return '?'; | |
| const whole = Math.max(0, Math.round(seconds)); const hours = Math.floor(whole / 3600); const minutes = Math.floor(whole % 3600 / 60); const remainder = whole % 60; | |
| return hours ? `${String(hours).padStart(2, '0')}:${String(minutes).padStart(2, '0')}:${String(remainder).padStart(2, '0')}` : `${String(minutes).padStart(2, '0')}:${String(remainder).padStart(2, '0')}`; | |
| } | |
| function resetProgress(kind) { progressStartedAt[kind] = null; setProgress(kind, 0); } | |
| function beginProgress(kind) { progressStartedAt[kind] = performance.now(); setProgress(kind, 0); } | |
| function setProgress(kind, completedSteps) { | |
| const completed = Math.max(0, Math.min(TOTAL_STEPS, Math.round(completedSteps))); const label = kind === 'forward' ? elements.forwardProgressLabel : elements.reverseProgressLabel; | |
| const elapsed = progressStartedAt[kind] === null ? 0 : (performance.now() - progressStartedAt[kind]) / 1000; const rate = elapsed > 0 && completed > 0 ? completed / elapsed : 0; const remaining = rate > 0 ? (TOTAL_STEPS - completed) / rate : Infinity; | |
| const text = `${completed}/${TOTAL_STEPS} [${formatDuration(elapsed)}<${completed === TOTAL_STEPS ? '00:00' : formatDuration(remaining)}, ${rate > 0 ? rate.toFixed(2) : '?'}it/s]`; label.value = text; label.textContent = text; | |
| } | |
| function stopAnimation() { if (animationTimer) window.clearInterval(animationTimer); animationTimer = null; document.querySelectorAll('.stage-card.is-active').forEach((card) => card.classList.remove('is-active')); } | |
| function runDiffusion() { | |
| try { computeForwardStates(); elements.configError.textContent = ''; } catch (error) { elements.configError.textContent = error.message; return; } | |
| stopAnimation(); forwardCompleted = false; forwardRevealCount = 1; beginProgress('forward'); renderTracks(); const cards = [...elements.diffusionTrack.querySelectorAll('.stage-card')]; cards[0].classList.add('is-active'); | |
| const duration = Math.max(1000, (checkpoints.length - 1) * 650); const startedAt = performance.now(); let activeIndex = 0; | |
| animationTimer = window.setInterval(() => { | |
| const completed = Math.min(TOTAL_STEPS, Math.floor((performance.now() - startedAt) / duration * TOTAL_STEPS)); setProgress('forward', completed); | |
| let nextIndex = activeIndex; while (nextIndex + 1 < checkpoints.length && checkpoints[nextIndex + 1] <= completed) nextIndex += 1; | |
| if (nextIndex !== activeIndex) { cards[activeIndex].classList.remove('is-active'); for (let index = activeIndex + 1; index <= nextIndex; index += 1) cards[index].classList.remove('is-pending'); cards[nextIndex].classList.add('is-active'); activeIndex = nextIndex; forwardRevealCount = nextIndex + 1; } | |
| if (completed >= TOTAL_STEPS) { forwardCompleted = true; setProgress('forward', TOTAL_STEPS); stopAnimation(); renderTracks(); } | |
| }, 50); | |
| } | |
| function parseTimesteps() { | |
| const values = elements.timestepsInput.value.split(',').map((part) => part.trim()).filter(Boolean).map(Number); | |
| if (values.length < 2 || values.some((value) => !Number.isInteger(value) || value < 0 || value > TOTAL_STEPS)) throw new Error('Enter at least two whole numbers between 0 and 1000.'); | |
| const unique = [...new Set(values)].sort((a, b) => a - b); if (unique[0] !== 0 || unique.at(-1) !== TOTAL_STEPS) throw new Error('Include both 0 and 1000 so each process has a start and end.'); if (unique.length > 12) throw new Error('Use no more than 12 timesteps.'); return unique; | |
| } | |
| function readSeed() { | |
| const value = Number(elements.seedInput.value.trim()); | |
| if (!Number.isSafeInteger(value) || value < 0 || value > 2147483647) throw new Error('Seed must be a whole number from 0 to 2147483647.'); | |
| return value; | |
| } | |
| function applyTimesteps() { | |
| try { const next = parseTimesteps(); stopAnimation(); if (generationController) generationController.abort(); checkpoints = next; reverseStates = new Map(); reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(readSeed()); forwardRevealCount = 1; forwardCompleted = false; computeForwardStates(); resetProgress('forward'); resetProgress('reverse'); elements.timestepsInput.value = checkpoints.join(', '); elements.configError.textContent = ''; setModelStatus('Timesteps updated. Select Denoise to run the model.'); renderTracks(); } | |
| catch (error) { elements.configError.textContent = error.message; } | |
| } | |
| function decodeFrame(data) { return new Promise((resolve, reject) => { const image = new Image(); image.onload = () => { const canvas = document.createElement('canvas'); canvas.width = image.naturalWidth; canvas.height = image.naturalHeight; canvas.getContext('2d').drawImage(image, 0, 0); resolve(canvas); }; image.onerror = () => reject(new Error('The backend returned an unreadable image frame.')); image.src = `data:image/jpeg;base64,${data}`; }); } | |
| async function generateReverse() { | |
| if (generationBusy) return; | |
| let seed; | |
| try { seed = readSeed(); elements.configError.textContent = ''; } catch (error) { elements.configError.textContent = error.message; return; } | |
| stopAnimation(); generationController = new AbortController(); generationBusy = true; reverseStates = new Map(); resetProgress('reverse'); setModelStatus('Loading model weights if needed, then sampling 1,000 reverse steps…'); renderTracks(); | |
| try { | |
| const initialImage = reverseInputMode === 'forward' ? stateToPng(reverseNoiseState) : null; | |
| const response = await fetch('/api/generate', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ model: modelInput.value, checkpoints, seed, initial_image: initialImage }), signal: generationController.signal }); if (!response.ok || !response.body) throw new Error(`Backend request failed (${response.status}).`); | |
| const reader = response.body.getReader(); const decoder = new TextDecoder(); let pending = ''; | |
| while (true) { const { value, done } = await reader.read(); pending += decoder.decode(value || new Uint8Array(), { stream: !done }); const lines = pending.split('\n'); pending = lines.pop(); for (const line of lines) if (line.trim()) await handleModelEvent(JSON.parse(line)); if (done) break; } | |
| if (pending.trim()) await handleModelEvent(JSON.parse(pending)); | |
| } catch (error) { if (error.name !== 'AbortError') setModelStatus(`Model error: ${error.message}`); } | |
| finally { generationBusy = false; generationController = null; renderTracks(); } | |
| } | |
| async function handleModelEvent(event) { | |
| if (event.type === 'status') setModelStatus(`Loading ${event.model} on the first generation…`); | |
| if (event.type === 'progress') { const completed = event.total - event.step; if (progressStartedAt.reverse === null) beginProgress('reverse'); setProgress('reverse', completed); setModelStatus(`Denoising · ${completed}/${event.total} steps complete.`); } | |
| if (event.type === 'frame') { reverseStates.set(event.step, await decodeFrame(event.image)); setProgress('reverse', TOTAL_STEPS - event.step); setModelStatus(`Real model inference · received t = ${event.step}.`); renderTracks(); const card = elements.denoisingTrack.querySelector(`[data-step="${event.step}"]`); if (card) { document.querySelectorAll('#denoising-track .is-active').forEach((item) => item.classList.remove('is-active')); card.classList.add('is-active'); } } | |
| if (event.type === 'complete') { setProgress('reverse', TOTAL_STEPS); setModelStatus(`Generated with ${event.model || modelInput.value} on ${event.device.toUpperCase()} · seed ${event.seed}.`); } if (event.type === 'error') throw new Error(event.message); | |
| } | |
| function loadImage(image, label) { const canvas = document.createElement('canvas'); canvas.width = SIZE; canvas.height = SIZE; const context = canvas.getContext('2d'); const scale = Math.max(SIZE / image.width, SIZE / image.height); const width = image.width * scale; const height = image.height * scale; context.drawImage(image, (SIZE - width) / 2, (SIZE - height) / 2, width, height); sourcePixels = canvasPixels(canvas); imageLabel = label; reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(readSeed()); forwardRevealCount = 1; forwardCompleted = false; resetProgress('forward'); computeForwardStates(); renderTracks(); } | |
| modelInput.addEventListener('change', () => { | |
| applyTimesteps(); | |
| const selectedModel = modelInput.value; | |
| const source = selectedModel === 'google/ddpm-cifar10-32' | |
| ? { path: 'assets/cifar10_train_23_32.png', label: 'CIFAR-10 train 23' } | |
| : { path: 'assets/celeba_hq_validation_91_256.png', label: 'Human face' }; | |
| const image = new Image(); | |
| image.onload = () => { | |
| if (modelInput.value === selectedModel) loadImage(image, source.label); | |
| }; | |
| image.onerror = () => setModelStatus('Could not load the selected model’s source image.'); | |
| image.src = source.path; | |
| }); | |
| elements.applyTimesteps.addEventListener('click', applyTimesteps); elements.timestepsInput.addEventListener('keydown', (event) => { if (event.key === 'Enter') applyTimesteps(); }); | |
| elements.randomizeSeed.addEventListener('click', () => { resampleNoise(); elements.configError.textContent = ''; forwardRevealCount = 1; forwardCompleted = false; resetProgress('forward'); computeForwardStates(); renderTracks(); }); | |
| elements.imageUpload.addEventListener('change', (event) => { const file = event.target.files?.[0]; if (!file) return; const image = new Image(); image.onload = () => { loadImage(image, file.name.replace(/\.[^.]+$/, '')); URL.revokeObjectURL(image.src); }; image.src = URL.createObjectURL(file); }); | |
| // Keep the original generated still life available while the default face loads, | |
| // and as a fallback if the image asset cannot be read. | |
| sourcePixels = canvasPixels(makeSourceArt()); computeForwardStates(); reverseNoiseState = makeNoiseState(readSeed()); resetProgress('forward'); resetProgress('reverse'); renderTracks(); | |
| const defaultSourceImage = new Image(); | |
| defaultSourceImage.onload = () => { | |
| if (modelInput.value === 'google/ddpm-celebahq-256') loadImage(defaultSourceImage, 'Human face'); | |
| }; | |
| defaultSourceImage.src = 'assets/celeba_hq_validation_91_256.png'; | |
| document.addEventListener('pointerdown', event => { | |
| if (!event.target.closest('.stage-canvas canvas')) document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible')); | |
| }); | |
| document.addEventListener('keydown', event => { | |
| if (event.key === 'Escape') document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible')); | |
| }); | |