Spaces:
Sleeping
Sleeping
File size: 22,327 Bytes
3030398 785e85b 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 785e85b 45769da 3030398 45769da 785e85b 45769da 785e85b 45769da 3030398 45769da 785e85b 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 3030398 45769da 785e85b 45769da 785e85b 45769da 785e85b 4f7577e 785e85b 45769da 0eff3a9 45769da 0eff3a9 785e85b 0eff3a9 785e85b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | 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'));
});
|