frontend/src/lib/App.svelte

<script lang="ts">
	import { PUBLIC_BACKEND_WS_URL } from '$env/static/public';
	import { onMount, tick } from 'svelte';
	import { nanoid } from 'nanoid';
	import { chatsStore, selectedChatId, loadingState, lastBase64Image } from '$lib/store';
	import type { Message, Chat, InferenceResponse, ImageFile } from '$lib/types';
	import { MessageType, Sender } from '$lib/types';
	import { timeFormater, fetchImageBase64 } from '$lib/utils';
	import ChatInput from '$lib/ChatInput.svelte';
	import ChatMessage from '$lib/ChatMessage.svelte';
	import ChatNewBtn from '$lib/ChatNewBtn.svelte';
	import IconDelete from './Icons/IconDelete.svelte';

	$: isLoading = !($loadingState === '' || $loadingState === 'Complete');

	let chatBoxEl: HTMLDivElement;
	let chatInputEl: HTMLInputElement;
	function clearStateMsg(t = 5000) {
		setTimeout(() => {
			$loadingState = '';
		}, t);
	}
	onMount(() => {
		const observer = new ResizeObserver(() => {
			window.scrollTo(0, chatBoxEl.getBoundingClientRect().height);
			if ('parentIFrame' in window) {
				(window as any).parentIFrame.scrollTo(0, chatBoxEl.getBoundingClientRect().height);
			}
		});
		observer.observe(chatBoxEl);
		// generateImage();
	});

	$: chatData = $chatsStore.find((chat) => chat.id === $selectedChatId);
	$: messages = chatData?.messages.sort((a, b) => a.timestamp - b.timestamp) || [];
	$: if ($selectedChatId) {
		// find last message with image type
		$lastBase64Image =
			[...messages].reverse().find((message) => message.type === MessageType.IMAGE)?.content ||
			null;
	}

	// this is when page starts or user deletes all chats
	function newChat() {
		const chatId = nanoid();
		const chat: Chat = {
			id: chatId,
			blurb: `New Chat - ${chatId}`,
			messages: [],
			timestamp: new Date().getTime()
		};
		$chatsStore = [chat].concat($chatsStore);
		$selectedChatId = chat.id;
	}
	function deleteChat(id: string) {
		$chatsStore = $chatsStore.filter((chat) => chat.id !== id);
	}

	// this is called when user send a text or image
	// if image then update last image store
	// if text and has image send to inference
	function submitMessage(event: CustomEvent) {
		if ($chatsStore.length === 0) {
			newChat();
		}
		const { type, content } = event.detail;
		const message: Message = {
			sender: Sender.USER,
			id: nanoid(),
			type: type,
			content: content,
			timestamp: new Date().getTime()
		};
		updateChatStore(message);

		if (type === MessageType.IMAGE) {
			// upate last image to be the last image sent
			$lastBase64Image = content;
		} else if (type === MessageType.TEXT) {
			// if the last message was an image, then we want to run inference
			// on the image and the text
			if (!$lastBase64Image) {
				const message: Message = {
					sender: Sender.BOT,
					id: nanoid(),
					type: MessageType.TEXT,
					content: "Sorry, I don't have an image to work with.",
					timestamp: new Date().getTime()
				};
				updateChatStore(message);
				return;
			}
			runInference($lastBase64Image, content);
		}
	}
	// hack to update store
	function updateChatStore(message: Message) {
		$chatsStore = $chatsStore.map((chat) => {
			if (chat.id === $selectedChatId) {
				chat.messages.push(message);
			}
			return chat;
		});
	}
	// run inference via websockets
	async function runInference(image: string, prompt: string) {
		if (isLoading || image === '' || prompt === '') {
			return;
		}

		$loadingState = 'Pending';

		const sessionHash = crypto.randomUUID();
		const hashpayload = {
			fn_index: 1,
			session_hash: sessionHash
		};
		const datapayload = {
			data: [
				prompt, // prompt
				10.5, // text guidance
				1.5, // image guidance
				image,
				15, // steps
				'', // negative promtp,
				512, // width
				512, // height
				0 // seed
			]
		};
		const websocket = new WebSocket(`wss://${PUBLIC_BACKEND_WS_URL}/queue/join`);
		// websocket.onopen = async function (event) {
		// 	websocket.send(JSON.stringify({ hash: sessionHash }));
		// };
		websocket.onclose = (evt) => {
			if (!evt.wasClean) {
				$loadingState = 'Error';
			}
		};
		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();

						return;
					case 'estimation':
						const { rank, queue_size } = data;
						$loadingState = `On queue ${rank}/${queue_size}`;
						break;
					case 'process_generating':
						$loadingState = data.success ? 'Generating' : 'Error';
						break;

					// here is success
					// got an image in response from inference
					// then update the chat store
					case 'process_completed':
						try {
							const response = data.output as InferenceResponse;
							const imageData = response.data[0] as ImageFile[];
							const nsfwData = data.output.data[1] as boolean[] | null;
							console.log('imageData', imageData);
							console.log('nsfwData', nsfwData);

							if (nsfwData && nsfwData[0]) {
								const message: Message = {
									sender: Sender.BOT,
									id: nanoid(),
									type: MessageType.TEXT,
									content:
										'Sorry this prompt possibly generates NSFW content. Please try another prompt.',
									timestamp: new Date().getTime()
								};
								updateChatStore(message);
							} else {
								const fileName = imageData[0].name;
								const imageBase64 = await fetchImageBase64(
									`https://${PUBLIC_BACKEND_WS_URL}/file=${fileName}`
								);
								const message: Message = {
									sender: Sender.BOT,
									id: nanoid(),
									type: MessageType.IMAGE,
									content: imageBase64,
									timestamp: new Date().getTime()
								};
								$lastBase64Image = imageBase64;
								updateChatStore(message);
							}
							$loadingState = data.success ? 'Complete' : 'Error';
							clearStateMsg();
						} catch (err) {
							const tError = err as Error;
							$loadingState = tError?.message;

							clearStateMsg(10000);
						}
						websocket.close();
						return;
					case 'process_starts':
						$loadingState = 'Processing';
						break;
				}
			} catch (e) {
				console.error(e);
				$loadingState = 'Error';
			}
		};
	}
</script>

<div>
	<h1 class="text-2xl">CHATS</h1>
	<div class="grid min-h-[40rem] grid-cols-4">
		<div class="col-span-1 flex flex-col border-r p-4 relative">
			<div class="sticky top-3">
				<ChatNewBtn on:click={newChat} />
				<div class="max-h-[40rem] flex flex-col gap-2 overflow-y-scroll">
					{#if $chatsStore.length}
						{#each $chatsStore as chat}
							<div class="flex flex-col relative">
								<button
									class="disabled:opacity-60 disabled:cursor-progress"
									on:click={() => ($selectedChatId = chat.id)}
									disabled={isLoading}
								>
									<div
										class=" flex flex-col h-16 items-start justify-center rounded-xl bg-gray-100 px-4 text-gray-900  
						{chat.id === $selectedChatId ? 'bg-gray-400' : ''}"
									>
										<h3 class="w-full truncate font-semibold">{chat.blurb}</h3>
										<p class="w-full truncate text-sm text-gray-500">
											{timeFormater(new Date(chat.timestamp))}
										</p>
									</div>
								</button>
								<button
									class="text-black absolute right-1 bottom-1 disabled:opacity-60"
									on:click={() => deleteChat(chat.id)}
									disabled={isLoading}
								>
									<IconDelete />
								</button>
							</div>
						{/each}
					{:else}
						<div
							class="flex h-16 flex-col items-start justify-center rounded-xl bg-gray-100 px-4 text-gray-900"
						>
							<h3 class="w-full truncate font-semibold">No chats</h3>
							<p class="w-full truncate text-sm text-gray-500">Start a new Chat!</p>
						</div>
					{/if}
				</div>
			</div>
		</div>

		<div class="col-span-3 flex flex-col" bind:this={chatBoxEl}>
			{#each messages as message}
				<ChatMessage {message} />
			{/each}
			<ChatInput on:submitMessage={submitMessage} bind:inputEl={chatInputEl} disabled={isLoading} />
		</div>
		<div class="top-0 right-0 z-10">
			Loading: {$loadingState}
		</div>
	</div>
</div>

<style lang="postcss" scoped>
</style>