Files
skymap-gen/src/main.js

699 lines
21 KiB
JavaScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import './style.css';
import * as THREE from 'three';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
import { convertFileSrc, invoke } from '@tauri-apps/api/tauri';
import { listen } from '@tauri-apps/api/event';
const app = document.getElementById('app');
app.innerHTML = `
<div id="canvas-container"></div>
<div id="thumb-dock">
<div id="thumb-header">Maps</div>
<div id="thumb-list"></div>
</div>
<div id="settings-panel" class="hidden">
<div class="settings-header">
<div>
<div class="settings-title">Generation Settings</div>
<div class="settings-sub">Defaults tuned for SDXL 360</div>
</div>
<button id="settings-close" aria-label="Close settings">×</button>
</div>
<div class="settings-grid">
<label>Steps
<input type="number" id="steps-input" min="1" max="80" />
</label>
<label>Guidance
<input type="number" id="guidance-input" step="0.1" min="0" max="20" />
</label>
<label>Width
<input type="number" id="width-input" min="256" max="4096" step="64" />
</label>
<label>Height
<input type="number" id="height-input" min="256" max="4096" step="64" />
</label>
<label>Scheduler
<select id="scheduler-input">
<option value="dpmsolver-sde">DPM++ 2M SDE</option>
<option value="dpmsolver">DPM++ 2M</option>
<option value="euler">Euler</option>
<option value="euler_a">Euler Ancestral</option>
<option value="heun">Heun</option>
<option value="ddim">DDIM</option>
</select>
</label>
<label>Upscale
<select id="upscale-input">
<option value="none">None</option>
<option value="realesrgan">Real-ESRGAN</option>
<option value="topaz">Topaz</option>
</select>
</label>
<label class="settings-check">Seam Inpainting
<input type="checkbox" id="seam-inpaint-input" />
</label>
<label>Model Path
<input type="text" id="model-path-input" />
</label>
<label>Base Model
<input type="text" id="base-model-input" />
</label>
<label>VAE
<input type="text" id="vae-model-input" />
</label>
</div>
<div class="settings-footer">
<button id="settings-reset">Reset defaults</button>
</div>
</div>
<div id="progress-overlay">
<div id="progress-ring">
<div id="progress-fill"></div>
<div id="progress-text">...</div>
</div>
</div>
<div id="prompt-bar">
<input id="prompt-input" type="text" placeholder="Describe the environment..." />
<button id="settings-btn" title="Generation settings">Settings</button>
<button id="generate-btn">Generate</button>
</div>
<div id="status"></div>
<div id="delete-confirm" class="hidden" role="dialog" aria-modal="true" aria-labelledby="delete-confirm-title">
<div id="delete-confirm-card">
<div id="delete-confirm-title">Delete map?</div>
<div id="delete-confirm-body">This will remove <span id="delete-confirm-filename"></span> from disk.</div>
<div id="delete-confirm-actions">
<button id="delete-confirm-cancel" type="button">Cancel</button>
<button id="delete-confirm-delete" type="button">Delete</button>
</div>
</div>
</div>
`;
const canvasContainer = document.getElementById('canvas-container');
const promptInput = document.getElementById('prompt-input');
const generateBtn = document.getElementById('generate-btn');
const settingsBtn = document.getElementById('settings-btn');
const statusEl = document.getElementById('status');
const thumbList = document.getElementById('thumb-list');
const thumbDock = document.getElementById('thumb-dock');
const deleteConfirm = document.getElementById('delete-confirm');
const deleteConfirmFilename = document.getElementById('delete-confirm-filename');
const deleteConfirmCancel = document.getElementById('delete-confirm-cancel');
const deleteConfirmDelete = document.getElementById('delete-confirm-delete');
const progressOverlay = document.getElementById('progress-overlay');
const progressFill = document.getElementById('progress-fill');
const progressText = document.getElementById('progress-text');
const settingsPanel = document.getElementById('settings-panel');
const settingsClose = document.getElementById('settings-close');
const settingsReset = document.getElementById('settings-reset');
const stepsInput = document.getElementById('steps-input');
const guidanceInput = document.getElementById('guidance-input');
const widthInput = document.getElementById('width-input');
const heightInput = document.getElementById('height-input');
const schedulerInput = document.getElementById('scheduler-input');
const upscaleInput = document.getElementById('upscale-input');
const seamInpaintInput = document.getElementById('seam-inpaint-input');
const modelPathInput = document.getElementById('model-path-input');
const baseModelInput = document.getElementById('base-model-input');
const vaeModelInput = document.getElementById('vae-model-input');
// Three.js setup
let renderer, scene, camera, controls, skyMesh;
const defaultTextureUrl = '/default.png';
const skyFadeDurationMs = 1000;
let skyTransition = null;
let skyTextureRequestId = 0;
let autoSpin = true;
let lastInteraction = Date.now();
const minZoomDistance = 0.2;
const maxZoomDistance = 20;
let desiredDistance = 1;
const zoomLerpFactor = 0.1;
const zoomVec = new THREE.Vector3();
let progressState = {
upscale: null,
seamInpaint: null,
phases: {
gen: null,
decode: null,
inpaint: null,
upscale: null,
},
};
let progressUnlisten = null;
let currentMapPath = null;
let deleteConfirmResolve = null;
let generationRunning = false;
let cancelRequested = false;
const defaultSettings = {
steps: 25,
width: 1536,
height: 768,
guidance: 6.5,
scheduler: 'dpmsolver-sde',
upscale: 'none',
seamInpaint: false,
modelPath: 'proximasan/sdxl-360-diffusion',
baseModel: 'stabilityai/stable-diffusion-xl-base-1.0',
vaeModel: 'madebyollin/sdxl-vae-fp16-fix',
};
let currentSettings = { ...defaultSettings };
function setStatus(msg) {
statusEl.textContent = msg || '';
statusEl.style.opacity = msg ? '0.98' : '0';
if (msg) {
clearTimeout(setStatus._timer);
setStatus._timer = setTimeout(() => {
statusEl.style.opacity = '0';
}, 5000);
}
}
async function loadTexture(url) {
const loader = new THREE.TextureLoader();
return new Promise((resolve, reject) => {
loader.load(
url,
tex => {
tex.mapping = THREE.EquirectangularReflectionMapping;
tex.colorSpace = THREE.SRGBColorSpace;
resolve(tex);
},
undefined,
err => reject(err)
);
});
}
async function setSkyFromUrl(url) {
const requestId = ++skyTextureRequestId;
try {
const texture = await loadTexture(url);
if (requestId !== skyTextureRequestId) {
texture.dispose();
return;
}
transitionSkyToTexture(texture);
} catch (e) {
console.error('Failed to load texture', e);
if (requestId === skyTextureRequestId) {
setStatus('Failed to load map texture');
}
}
}
function replaceSkyTexture(texture) {
const previousTexture = skyMesh.material.map;
skyMesh.material.map = texture;
skyMesh.material.needsUpdate = true;
if (previousTexture && previousTexture !== texture) {
previousTexture.dispose();
}
}
function removeSkyTransition(transition) {
scene.remove(transition.mesh);
transition.material.dispose();
}
function finishSkyTransition(transition) {
if (!transition || skyTransition !== transition) return;
removeSkyTransition(transition);
skyTransition = null;
replaceSkyTexture(transition.texture);
}
function transitionSkyToTexture(texture) {
if (!skyMesh.material.map) {
replaceSkyTexture(texture);
return;
}
if (skyTransition) {
finishSkyTransition(skyTransition);
}
const material = new THREE.MeshBasicMaterial({
side: THREE.BackSide,
color: 0xffffff,
map: texture,
transparent: true,
opacity: 0,
depthWrite: false,
});
const mesh = new THREE.Mesh(skyMesh.geometry, material);
mesh.rotation.copy(skyMesh.rotation);
mesh.renderOrder = skyMesh.renderOrder + 1;
scene.add(mesh);
skyTransition = {
mesh,
material,
texture,
start: performance.now(),
duration: skyFadeDurationMs,
};
}
function updateSkyTransition(now) {
if (!skyTransition) return;
const progress = Math.min(1, (now - skyTransition.start) / skyTransition.duration);
skyTransition.material.opacity = THREE.MathUtils.smoothstep(progress, 0, 1);
skyTransition.mesh.rotation.copy(skyMesh.rotation);
if (progress >= 1) {
finishSkyTransition(skyTransition);
}
}
function initScene() {
scene = new THREE.Scene();
camera = new THREE.PerspectiveCamera(75, window.innerWidth / window.innerHeight, 0.1, 1000);
camera.position.set(0.1, 0, 0.1);
renderer = new THREE.WebGLRenderer({ antialias: true });
renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
renderer.setSize(window.innerWidth, window.innerHeight);
renderer.outputColorSpace = THREE.SRGBColorSpace;
renderer.toneMapping = THREE.ACESFilmicToneMapping;
canvasContainer.appendChild(renderer.domElement);
desiredDistance = camera.position.length();
const geometry = new THREE.SphereGeometry(50, 64, 32);
const material = new THREE.MeshBasicMaterial({ side: THREE.BackSide, color: 0xffffff });
skyMesh = new THREE.Mesh(geometry, material);
scene.add(skyMesh);
controls = new OrbitControls(camera, renderer.domElement);
controls.enablePan = false;
controls.enableZoom = false; // we implement custom smooth zooming
controls.enableDamping = true;
controls.dampingFactor = 0.08;
controls.rotateSpeed = 0.3;
controls.autoRotate = false;
controls.autoRotateSpeed = 0.15;
controls.minDistance = 0.05;
controls.maxDistance = 20;
controls.addEventListener('start', () => {
autoSpin = false;
lastInteraction = Date.now();
});
controls.addEventListener('end', () => {
lastInteraction = Date.now();
});
window.addEventListener('resize', onResize);
renderer.domElement.addEventListener(
'wheel',
(e) => {
e.preventDefault();
const delta = e.deltaY;
// Scale desired distance exponentially for smooth feel
const factor = Math.exp(delta * 0.0015);
desiredDistance = THREE.MathUtils.clamp(desiredDistance * factor, minZoomDistance, maxZoomDistance);
lastInteraction = Date.now();
autoSpin = false;
},
{ passive: false }
);
animate();
}
function onResize() {
if (!renderer || !camera) return;
const w = window.innerWidth;
const h = window.innerHeight;
renderer.setSize(w, h);
camera.aspect = w / h;
camera.updateProjectionMatrix();
}
function animate() {
requestAnimationFrame(animate);
const now = performance.now();
const idle = Date.now() - lastInteraction > 2000;
if (idle) autoSpin = true;
if (autoSpin) {
skyMesh.rotation.y += 0.0008;
}
updateSkyTransition(now);
// Smooth zoom toward desired distance
const target = controls.target;
const currentDistance = camera.position.distanceTo(target);
const nextDistance = THREE.MathUtils.lerp(currentDistance, desiredDistance, zoomLerpFactor);
zoomVec.copy(camera.position).sub(target).normalize().multiplyScalar(nextDistance).add(target);
camera.position.copy(zoomVec);
controls.update();
renderer.render(scene, camera);
}
async function pathToObjectUrl(path) {
return convertFileSrc(path);
}
function resetProgressState() {
progressState = {
upscale: null,
seamInpaint: null,
phases: {
gen: null,
decode: null,
inpaint: null,
upscale: null,
},
};
}
function computeProgress() {
const upscaleOn = progressState.upscale === true;
const seamOn = progressState.seamInpaint === true;
const weights = {
gen: seamOn || upscaleOn ? 0.45 : 0.9,
decode: seamOn || upscaleOn ? 0.05 : 0.1,
inpaint: seamOn ? 0.5 : 0,
upscale: upscaleOn ? 0.5 : 0,
};
if (seamOn && upscaleOn) {
weights.gen = 0.25;
weights.decode = 0.05;
weights.inpaint = 0.25;
weights.upscale = 0.45;
}
const frac = phase => {
if (!phase || !phase.total) return 0;
const pct = phase.current / phase.total;
return Math.max(0, Math.min(1, pct));
};
return (
frac(progressState.phases.gen) * weights.gen +
frac(progressState.phases.decode) * weights.decode +
frac(progressState.phases.inpaint) * weights.inpaint +
frac(progressState.phases.upscale) * weights.upscale
);
}
function updateProgressDisplay(pct) {
if (!progressOverlay) return;
const clamped = Math.max(0, Math.min(1, pct || 0));
const deg = clamped * 100;
progressFill.style.background = `conic-gradient(#ffffff ${deg}%, rgba(255,255,255,0.12) ${deg}%)`;
progressText.textContent = `${Math.round(clamped * 100)}%`;
}
function startProgress() {
resetProgressState();
if (progressOverlay) {
progressOverlay.style.display = 'flex';
}
updateProgressDisplay(0);
}
function stopProgress() {
updateProgressDisplay(1);
if (progressOverlay) {
progressOverlay.style.display = 'none';
}
}
async function refreshThumbnails(selectedPath, shouldLoadSelection = true) {
let maps;
try {
maps = await invoke('list_maps');
} catch (e) {
console.error(e);
setStatus('Failed to read map list');
thumbDock.style.display = 'none';
return;
}
thumbList.innerHTML = '';
if (!maps || maps.length === 0) {
thumbDock.style.display = 'none';
return;
}
thumbDock.style.display = 'block';
for (let idx = 0; idx < maps.length; idx++) {
const item = maps[idx];
const el = document.createElement('div');
el.className = 'thumb-item';
const img = document.createElement('img');
const deleteBtn = document.createElement('button');
deleteBtn.className = 'thumb-delete';
deleteBtn.type = 'button';
deleteBtn.textContent = '×';
deleteBtn.title = 'Delete map';
deleteBtn.setAttribute('aria-label', `Delete ${item.filename}`);
deleteBtn.addEventListener('click', async (event) => {
event.preventDefault();
event.stopPropagation();
await deleteMap(item, el);
});
let fileUrl;
try {
fileUrl = await pathToObjectUrl(item.path);
img.src = fileUrl;
} catch (err) {
console.error('Failed to load thumbnail', err);
img.alt = 'Failed to load';
}
el.appendChild(img);
el.appendChild(deleteBtn);
el.title = item.filename;
el.addEventListener('click', () => {
if (fileUrl) {
currentMapPath = item.path;
setSkyFromUrl(fileUrl);
setStatus(`Loaded ${item.filename}`);
}
});
thumbList.appendChild(el);
if (fileUrl) {
if (shouldLoadSelection && idx === 0 && !selectedPath) {
currentMapPath = item.path;
setSkyFromUrl(fileUrl);
setStatus(`Showing ${item.filename}`);
}
if (shouldLoadSelection && selectedPath && selectedPath === item.path) {
currentMapPath = item.path;
setSkyFromUrl(fileUrl);
setStatus(`Showing ${item.filename}`);
}
}
}
}
async function deleteMap(item, tileEl) {
if (!item?.path) return;
const confirmed = await requestDeleteConfirmation(item.filename);
if (!confirmed) return;
try {
await invoke('delete_map', { path: item.path });
const wasCurrent = currentMapPath === item.path;
tileEl?.remove();
setStatus(`Deleted ${item.filename}`);
if (thumbList.children.length === 0) {
thumbDock.style.display = 'none';
currentMapPath = null;
await setSkyFromUrl(defaultTextureUrl);
return;
}
if (wasCurrent) {
thumbList.querySelector('.thumb-item')?.click();
}
} catch (err) {
console.error(err);
setStatus('Failed to delete map');
}
}
function requestDeleteConfirmation(filename) {
if (deleteConfirmResolve) {
closeDeleteConfirmation(false);
}
deleteConfirmFilename.textContent = `"${filename}"`;
deleteConfirm.classList.remove('hidden');
deleteConfirmDelete.focus();
return new Promise(resolve => {
deleteConfirmResolve = resolve;
});
}
function closeDeleteConfirmation(result) {
deleteConfirm.classList.add('hidden');
if (deleteConfirmResolve) {
const resolve = deleteConfirmResolve;
deleteConfirmResolve = null;
resolve(result);
}
}
async function generateMap() {
if (generationRunning) {
await cancelGeneration();
return;
}
const prompt = promptInput.value.trim();
if (!prompt) {
setStatus('Please enter a prompt');
return;
}
const settings = {
steps: Number(stepsInput.value) || defaultSettings.steps,
guidance: Number(guidanceInput.value) || defaultSettings.guidance,
width: Number(widthInput.value) || defaultSettings.width,
height: Number(heightInput.value) || defaultSettings.height,
scheduler: schedulerInput.value || defaultSettings.scheduler,
upscale: upscaleInput.value || defaultSettings.upscale,
seamInpaint: seamInpaintInput.checked,
modelPath: modelPathInput.value.trim() || defaultSettings.modelPath,
baseModel: baseModelInput.value.trim() || defaultSettings.baseModel,
vaeModel: vaeModelInput.value.trim() || defaultSettings.vaeModel,
};
currentSettings = settings;
generationRunning = true;
cancelRequested = false;
generateBtn.disabled = false;
generateBtn.textContent = 'Cancel';
setStatus('Generating...');
startProgress();
try {
const result = await invoke('generate_map', { prompt, settings });
const outputPath = result.outputPath || result.output_path || result;
if (outputPath) {
const fileUrl = await pathToObjectUrl(outputPath);
await refreshThumbnails(outputPath, false);
currentMapPath = outputPath;
await setSkyFromUrl(fileUrl);
setStatus('New environment loaded');
} else {
setStatus('Generation finished, but no output path reported');
}
} catch (e) {
console.error(e);
setStatus(e === 'Generation cancelled' ? 'Generation cancelled' : (typeof e === 'string' ? e : 'Generation failed'));
} finally {
stopProgress();
generationRunning = false;
cancelRequested = false;
generateBtn.textContent = 'Generate';
generateBtn.disabled = false;
}
}
async function cancelGeneration() {
if (!generationRunning || cancelRequested) return;
cancelRequested = true;
generateBtn.disabled = true;
generateBtn.textContent = 'Cancelling...';
setStatus('Cancelling generation...');
try {
await invoke('cancel_generation');
} catch (e) {
console.error(e);
cancelRequested = false;
generateBtn.disabled = false;
generateBtn.textContent = 'Cancel';
setStatus(typeof e === 'string' ? e : 'Failed to cancel generation');
}
}
function setupUI() {
generateBtn.addEventListener('click', () => {
if (generationRunning) {
cancelGeneration();
} else {
generateMap();
}
});
promptInput.addEventListener('keydown', (e) => {
if (e.key === 'Enter' && !generationRunning) generateMap();
});
settingsBtn.addEventListener('click', () => {
settingsPanel.classList.remove('hidden');
});
settingsClose.addEventListener('click', () => {
settingsPanel.classList.add('hidden');
});
settingsReset.addEventListener('click', () => {
applySettings(defaultSettings);
});
deleteConfirmCancel.addEventListener('click', () => {
closeDeleteConfirmation(false);
});
deleteConfirmDelete.addEventListener('click', () => {
closeDeleteConfirmation(true);
});
deleteConfirm.addEventListener('click', (event) => {
if (event.target === deleteConfirm) {
closeDeleteConfirmation(false);
}
});
document.addEventListener('keydown', (event) => {
if (event.key === 'Escape' && !deleteConfirm.classList.contains('hidden')) {
closeDeleteConfirmation(false);
}
});
}
async function setupProgressEvents() {
try {
progressUnlisten = await listen('gen-progress', (event) => {
const data = event?.payload || {};
if (typeof data.upscale === 'boolean') {
progressState.upscale = data.upscale;
}
if (typeof data.seamInpaint === 'boolean') {
progressState.seamInpaint = data.seamInpaint;
}
if (data.phase && typeof data.current === 'number' && typeof data.total === 'number') {
progressState.phases[data.phase] = {
current: data.current,
total: data.total,
};
}
updateProgressDisplay(computeProgress());
});
} catch (err) {
console.error('Failed to bind progress listener', err);
}
}
async function bootstrap() {
initScene();
setupUI();
applySettings(defaultSettings);
await setupProgressEvents();
await setSkyFromUrl(defaultTextureUrl);
await refreshThumbnails();
setStatus('Ready');
}
function applySettings(cfg) {
currentSettings = { ...cfg };
stepsInput.value = cfg.steps;
guidanceInput.value = cfg.guidance;
widthInput.value = cfg.width;
heightInput.value = cfg.height;
schedulerInput.value = cfg.scheduler;
upscaleInput.value = cfg.upscale;
seamInpaintInput.checked = Boolean(cfg.seamInpaint);
modelPathInput.value = cfg.modelPath;
baseModelInput.value = cfg.baseModel;
vaeModelInput.value = cfg.vaeModel;
}
bootstrap();