frontend/src/routes/+page.svelte

<script lang="ts">
	import { onMount } from 'svelte';
	import type { ColorsPrompt, ColorsImage } from '$lib/types';
	import { randomSeed, extractPalette, uploadImage } from '$lib/utils';
	import { isLoading, loadingState } from '$lib/store';
	import { PUBLIC_WS_ENDPOINT, PUBLIC_API} from '$env/static/public';
	import Pallette from '$lib/Palette.svelte';
	import ArrowRight from '$lib/ArrowRight.svelte';
	import ArrowLeft from '$lib/ArrowLeft.svelte';

	let promptsData: ColorsPrompt[] = [];
	let prompt: string;
	let promptInputEl: HTMLElement;

	onMount(() => {
		fetchData();
		const interval = window.setInterval(fetchData, 5000);
		return () => {
			clearInterval(interval);
		};
	});

	async function fetchData() {
		const palettes = await fetch(PUBLIC_API + '/data').then((d) => d.json());
		if (!promptsData || palettes?.length > promptsData?.length) {
			promptsData = sortData(palettes);
		}
	}

	$: promptsTotal = promptsData?.length || null;

	let page: number = 0;
	const maxPerPage: number = 10;
	$: totalPages = Math.ceil(promptsData?.length / maxPerPage) || 0;
	$: promptsDataPage = [...promptsData].slice(page * maxPerPage, (page + 1) * maxPerPage);
	let pagesLinks: number[] = [];
	$: if (totalPages) {
		const pagesNums = Array(totalPages)
			.fill([])
			.map((_, i) => ({ value: i, label: i + 1 }));
		pagesLinks = pagesNums
			.slice(0, 3)
			.concat([{ value: -1, label: '...' }])
			.concat(pagesNums.length > 3 ? pagesNums.slice(-1) : []);
		console.log(pagesLinks);
	}

	function sortData(_promptData: ColorsPrompt[]) {
		return _promptData
			.sort((a, b) => b.id - a.id)
			.map((p) => p.data)
			.filter((d) => d.images.length > 0);
	}
	async function savePaletteDB(colorPrompt: ColorsPrompt) {
		try {
			const newPalettes: ColorsPrompt[] = await fetch(PUBLIC_API + '/new_palette', {
				method: 'POST',
				headers: {
					'Content-Type': 'application/json'
				},
				body: JSON.stringify({
					prompt: colorPrompt.prompt,
					images: colorPrompt.images.map((i) => ({
						imgURL: i.imgURL,
						colors: i.colors.map((c) => c.formatHex())
					}))
				})
			}).then((d) => d.json());

			promptsData = sortData(newPalettes);
		} catch (e) {
			console.error(e);
		}
	}

	async function generatePalette(_prompt: string) {
		if (!_prompt || $isLoading == true) return;
		$loadingState = 'Pending';
		$isLoading = true;

		const sessionHash = crypto.randomUUID();

		const hashpayload = {
			fn_index: 2,
			session_hash: sessionHash
		};

		const datapayload = {
			data: [_prompt]
		};

		const websocket = new WebSocket(PUBLIC_WS_ENDPOINT);
		// websocket.onopen = async function (event) {
		// 	websocket.send(JSON.stringify({ hash: sessionHash }));
		// };
		websocket.onclose = (evt) => {
			if (!evt.wasClean) {
				$loadingState = 'Error';
				$isLoading = false;
			}
		};
		websocket.onmessage = async function (event) {
			try {
				const data = JSON.parse(event.data);
				$loadingState = '';
				switch (data.msg) {
					case 'send_hash':
						websocket.send(JSON.stringify(hashpayload));
						break;
					case 'send_data':
						$loadingState = 'Sending Data';
						websocket.send(JSON.stringify({ ...hashpayload, ...datapayload }));
						break;
					case 'queue_full':
						$loadingState = 'Queue full';
						websocket.close();
						$isLoading = false;
						return;
					case 'estimation':
						const { msg, rank, queue_size } = data;
						$loadingState = `On queue ${rank}/${queue_size}`;
						break;
					case 'process_generating':
						$loadingState = data.success ? 'Generating' : 'Error';
						break;
					case 'process_completed':
						try {
							const images = await extractColorsImages(data.output.data[0], _prompt);
							savePaletteDB({
								prompt: _prompt,
								images
							});
							$loadingState = data.success ? 'Complete' : 'Error';
						} catch (e) {
							$loadingState = e.message;
						}
						websocket.close();
						$isLoading = false;
						return;
					case 'process_starts':
						$loadingState = 'Processing';
						break;
				}
			} catch (e) {
				console.error(e);
				$isLoading = false;
				$loadingState = 'Error';
			}
		};
	}
	async function extractColorsImages(images: string[], _prompt: string): Promise<ColorsImage[]> {
		const nsfwColors = ['#040404', '#B7B7B7', '#565656', '#747474', '#6C6C6C'];

		const colorImages = [];
		let isNSFW = false;
		for (const base64img of images) {
			const { colors, imgBlob } = await extractPalette(base64img);
			if (
				!colors.map((color) => color.formatHex().toUpperCase()).every((c) => nsfwColors.includes(c))
			) {
				const url = await uploadImage(imgBlob, _prompt);
				const colorsImage: ColorsImage = {
					colors,
					imgURL: url
				};
				colorImages.push(colorsImage);
			} else {
				isNSFW = true;
			}
		}

		if (colorImages.length === 0 && isNSFW) {
			console.error('Possible NSFW image');
			throw new Error('Possible NSFW image');
		}
		return colorImages;
	}
	function remix(e: CustomEvent) {
		prompt = e.detail.prompt;
		promptInputEl.scrollIntoView({ behavior: 'smooth' });
		scrollTop();
	}
	function scrollTop() {
		window.scrollTo(0, 0);
		if ('parentIFrame' in window) {
			window.parentIFrame.scrollTo(0, promptInputEl.offsetTop);
		}
	}
</script>

<div class="max-w-screen-md mx-auto px-3 py-8 relative z-0">
	<h1 class="text-3xl font-bold leading-normal">Palette generation with Stable Diffusion</h1>
	<p class="text-sm">
		Original ideas:

		<a
			class="link"
			target="_blank"
			rel="nofollow noopener"
			href="https://twitter.com/mattdesl/status/1569457653298139136"
		>
			Matt DesLauriers
		</a>,
		<a class="link" href="https://drib.net/homage"> dribnet </a>
	</p>
	<div class="relative top-0 z-50 bg-white dark:bg-black py-3">
		<form class="grid grid-cols-6" on:submit|preventDefault={() => generatePalette(prompt)}>
			<input
				bind:this={promptInputEl}
				class="input"
				placeholder="A photo of a beautiful sunset in San Francisco"
				title="Input prompt to generate image and obtain palette"
				type="text"
				name="prompt"
				bind:value={prompt}
				disabled={$isLoading}
			/>
			<button
				class="button"
				on:click|preventDefault={() => generatePalette(prompt)}
				disabled={$isLoading}
				title="Generate Palette"
			>
				Create Palette
			</button>
		</form>
		{#if $loadingState}
			<h3 class="text-xs font-bold ml-3 inline-block">{$loadingState}</h3>
			{#if $isLoading}
				<svg
					xmlns="http://www.w3.org/2000/svg"
					fill="none"
					viewBox="0 0 24 24"
					class="animate-spin max-w-[1rem] inline-block"
				>
					<path
						fill="currentColor"
						d="M20 12a8 8 0 0 1-8 8v4a12 12 0 0 0 12-12h-4Zm-2-5.3a8 8 0 0 1 2 5.3h4c0-3-1.1-5.8-3-8l-3 2.7Z"
					/>
				</svg>
			{/if}
		{/if}
	</div>

	<div class="flex items-center gap-4 my-10">
		<div class="font-bold text-sm">
			{promptsTotal ? `${promptsTotal} submitted palettes` : 'Loading...'}
		</div>
		<div class="grow border-b border-gray-200" />
	</div>

	{#if promptsDataPage}
		<div>
			{#each promptsDataPage as promptData}
				<Pallette {promptData} on:remix={remix} />
				<div class="border-b border-gray-200 py-2" />
			{/each}
		</div>
		<nav role="navigation">
			<ul
				class="items-center sm:justify-center space-x-2  select-none w-full flex justify-center mt-6 mb-4"
			>
				<li />
				<li>
					<a
						on:click|preventDefault={() => {
							page = page - 1 < 0 ? 0 : page - 1;
							scrollTop();
						}}
						class="px-2.5 py-1 hover:bg-gray-50 dark:hover:bg-gray-800 flex items-center rounded-lg"
						href="#"
						><ArrowLeft /> Previous
					</a>
				</li>
				<li class="text-sm">
					<span class="inline-block min-w-[3ch] text-right">{page + 1} </span>/<span
						class="inline-block min-w-[3ch]"
						>{totalPages}
					</span>
				</li>
				<li>
					<a
						on:click|preventDefault={() => {
							page = page + 1 >= totalPages - 1 ? totalPages - 1 : page + 1;
							scrollTop();
						}}
						class="px-2.5 py-1 hover:bg-gray-50 dark:hover:bg-gray-800 flex items-center rounded-lg"
						href="#"
						>Next <ArrowRight />
					</a>
				</li>
			</ul>
		</nav>
	{/if}
</div>

<style lang="postcss" scoped>
	.link {
		@apply text-xs underline font-bold hover:no-underline hover:text-gray-500 visited:text-gray-500;
	}
	.input {
		@apply text-sm disabled:opacity-50 col-span-4 md:col-span-5 italic dark:placeholder:text-black placeholder:text-white text-white dark:text-black placeholder:text-opacity-30 dark:placeholder:text-opacity-10 dark:bg-white bg-slate-900 border-2 border-black rounded-2xl px-2 shadow-sm focus:outline-none focus:border-gray-400 focus:ring-1;
	}
	.button {
		@apply disabled:opacity-50 col-span-2 md:col-span-1 dark:bg-white dark:text-black border-2 border-black rounded-2xl ml-2 px-2 py-2  text-xs shadow-sm font-bold focus:outline-none focus:border-gray-400;
	}
</style>