[Assistants] Use textToImage task for avatar generation (#662)

* Generate assistants avatar using stablediffusion

* wording

* Update +page.server.ts

Co-authored-by: Michael Fried <mikelfried@gmail.com>

* Add timeout & controls to avatar generation

* Add controls for avatar generation in .env

* Update src/routes/+layout.server.ts

Co-authored-by: Mishig <mishig.davaadorj@coloradocollege.edu>

* Update src/lib/components/AssistantSettings.svelte

Co-authored-by: Mishig <mishig.davaadorj@coloradocollege.edu>

* Fix avatar gen feature flag

* Can only upload avatar if generate is unchecked

---------

Co-authored-by: Michael Fried <mikelfried@gmail.com>
Co-authored-by: Mishig <mishig.davaadorj@coloradocollege.edu>

Author: 3 people

.env

@@ -128,4 +128,6 @@ LLM_SUMMERIZATION=true
 # PUBLIC_APP_DATA_SHARING=1
 # PUBLIC_APP_DISCLAIMER=1
 
-ENABLE_ASSISTANTS=false #set to true to enable assistants feature
+ENABLE_ASSISTANTS=false #set to true to enable assistants feature
+ASSISTANTS_GENERATE_AVATAR=true #requires an hf token, uses the model description and name to generate an avatar using a text to image model
+TEXT_TO_IMAGE_MODEL="runwayml/stable-diffusion-v1-5"

src/lib/components/AssistantSettings.svelte

@@ -4,10 +4,12 @@
 	import type { Assistant } from "$lib/types/Assistant";
 
 	import { onMount } from "svelte";
-	import { enhance } from "$app/forms";
+	import { applyAction, enhance } from "$app/forms";
 	import { base } from "$app/paths";
 	import CarbonPen from "~icons/carbon/pen";
 	import { useSettingsStore } from "$lib/stores/settings";
+	import { page } from "$app/stores";
+	import IconLoading from "./icons/IconLoading.svelte";
 
 	type ActionData = {
 		error: boolean;
@@ -49,13 +51,18 @@
 	function getError(field: string, returnForm: ActionData) {
 		return returnForm?.errors.find((error) => error.field === field)?.message ?? "";
 	}
+
+	let loading = false;
+
+	let generateAvatar = false;
 </script>
 
 <form
 	method="POST"
 	class="h-full w-full overflow-x-clip"
 	enctype="multipart/form-data"
 	use:enhance={async ({ formData }) => {
+		loading = true;
 		const avatar = formData.get("avatar");
 
 		if (avatar && typeof avatar !== "string" && avatar.size > 0 && compress) {
@@ -67,6 +74,11 @@
 				formData.set("avatar", resizedImage);
 			});
 		}
+
+		return async ({ result }) => {
+			loading = false;
+			await applyAction(result);
+		};
 	}}
 >
 	{#if assistant}
@@ -90,6 +102,7 @@
 					accept="image/*"
 					name="avatar"
 					class="invisible z-10 block h-0 w-0"
+					disabled={generateAvatar}
 					on:change={onFilesChange}
 				/>
 				{#if (files && files[0]) || assistant?.avatar}
@@ -122,9 +135,24 @@
 						Reset
 					</button>
 				{:else}
-					<span class="text-xs text-gray-500 hover:underline">Click to upload</span>
+					<span
+						class="text-xs text-gray-500"
+						class:hover:underline={!generateAvatar}
+						class:cursor-pointer={!generateAvatar}>Click to upload</span
+					>
 				{/if}
 				<p class="text-xs text-red-500">{getError("avatar", form)}</p>
+				{#if !files?.[0] && $page.data.avatarGeneration && !assistant?.avatar}
+					<label class="text-xs text-gray-500">
+						<input
+							type="checkbox"
+							name="generateAvatar"
+							class="text-xs text-gray-500"
+							bind:checked={generateAvatar}
+						/>
+						Generate avatar from description
+					</label>
+				{/if}
 			</label>
 
 			<label>
@@ -220,8 +248,19 @@
 			class="rounded-full bg-gray-200 px-8 py-2 font-semibold text-gray-600">Cancel</a
 		>
 
-		<button type="submit" class="rounded-full bg-black px-8 py-2 font-semibold text-white md:px-20"
-			>{assistant ? "Save" : "Create"}</button
+		<button
+			type="submit"
+			disabled={loading}
+			aria-disabled={loading}
+			class="rounded-full bg-black px-8 py-2 font-semibold md:px-20"
+			class:bg-gray-200={loading}
+			class:text-gray-600={loading}
+			class:text-white={!loading}
 		>
+			{assistant ? "Save" : "Create"}
+			{#if loading}
+				<IconLoading classNames="ml-2 h-min" />
+			{/if}
+		</button>
 	</div>
 </form>

src/lib/utils/generateAvatar.ts

@@ -0,0 +1,24 @@
+import { HF_TOKEN, TEXT_TO_IMAGE_MODEL } from "$env/static/private";
+import { generateFromDefaultEndpoint } from "$lib/server/generateFromDefaultEndpoint";
+import { HfInference } from "@huggingface/inference";
+
+export async function generateAvatar(description?: string, name?: string): Promise<File> {
+	const queryPrompt = `Generate a prompt for an image-generation model for the following: 
+Name: ${name}
+Description: ${description}
+`;
+	const imagePrompt = await generateFromDefaultEndpoint({
+		messages: [{ from: "user", content: queryPrompt }],
+		preprompt:
+			"You are an assistant tasked with generating simple image descriptions. The user will ask you for an image, based on the name and a description of what they want, and you should reply with a short, concise, safe, descriptive sentence.",
+	});
+
+	const hf = new HfInference(HF_TOKEN);
+
+	const blob = await hf.textToImage({
+		inputs: imagePrompt,
+		model: TEXT_TO_IMAGE_MODEL,
+	});
+
+	return new File([blob], "avatar.png");
+}

src/lib/utils/timeout.ts

@@ -0,0 +1,6 @@
+export const timeout = <T>(prom: Promise<T>, time: number): Promise<T> => {
+	let timer: NodeJS.Timeout;
+	return Promise.race([prom, new Promise<T>((_r, rej) => (timer = setTimeout(rej, time)))]).finally(
+		() => clearTimeout(timer)
+	);
+};

src/routes/+layout.server.ts

@@ -13,6 +13,8 @@ import {
 	YDC_API_KEY,
 	USE_LOCAL_WEBSEARCH,
 	ENABLE_ASSISTANTS,
+	ASSISTANTS_GENERATE_AVATAR,
+	TEXT_TO_IMAGE_MODEL,
 } from "$env/static/private";
 import { ObjectId } from "mongodb";
 import type { ConvSidebar } from "$lib/types/ConvSidebar";
@@ -161,6 +163,7 @@ export const load: LayoutServerLoad = async ({ locals, depends }) => {
 			email: locals.user.email,
 		},
 		assistant,
+		avatarGeneration: ASSISTANTS_GENERATE_AVATAR === "true" && !!TEXT_TO_IMAGE_MODEL,
 		enableAssistants,
 		loginRequired,
 		loginEnabled: requiresUser,

src/routes/settings/assistants/[assistantId]/edit/+page.server.ts

@@ -7,6 +7,9 @@ import { ObjectId } from "mongodb";
 import { z } from "zod";
 import sizeof from "image-size";
 import { sha256 } from "$lib/utils/sha256";
+import { ASSISTANTS_GENERATE_AVATAR, HF_TOKEN } from "$env/static/private";
+import { generateAvatar } from "$lib/utils/generateAvatar";
+import { timeout } from "$lib/utils/timeout";
 
 const newAsssistantSchema = z.object({
 	name: z.string().min(1),
@@ -18,6 +21,10 @@ const newAsssistantSchema = z.object({
 	exampleInput3: z.string().optional(),
 	exampleInput4: z.string().optional(),
 	avatar: z.instanceof(File).optional(),
+	generateAvatar: z
+		.literal("on")
+		.optional()
+		.transform((el) => !!el),
 });
 
 const uploadAvatar = async (avatar: File, assistantId: ObjectId): Promise<string> => {
@@ -99,6 +106,29 @@ export const actions: Actions = {
 			}
 
 			hash = await uploadAvatar(parse.data.avatar, assistant._id);
+		} else if (
+			ASSISTANTS_GENERATE_AVATAR === "true" &&
+			HF_TOKEN !== "" &&
+			parse.data.generateAvatar
+		) {
+			try {
+				const avatar = await timeout(
+					generateAvatar(parse.data.description, parse.data.name),
+					30000
+				);
+
+				hash = await uploadAvatar(avatar, assistant._id);
+			} catch (err) {
+				return fail(400, {
+					error: true,
+					errors: [
+						{
+							field: "avatar",
+							message: "Avatar generation failed. Try again or disable the feature.",
+						},
+					],
+				});
+			}
 		}
 
 		const { acknowledged } = await collections.assistants.replaceOne(

src/routes/settings/assistants/new/+page.server.ts

@@ -7,6 +7,9 @@ import { ObjectId } from "mongodb";
 import { z } from "zod";
 import sizeof from "image-size";
 import { sha256 } from "$lib/utils/sha256";
+import { ASSISTANTS_GENERATE_AVATAR, HF_TOKEN } from "$env/static/private";
+import { timeout } from "$lib/utils/timeout";
+import { generateAvatar } from "$lib/utils/generateAvatar";
 
 const newAsssistantSchema = z.object({
 	name: z.string().min(1),
@@ -18,6 +21,10 @@ const newAsssistantSchema = z.object({
 	exampleInput3: z.string().optional(),
 	exampleInput4: z.string().optional(),
 	avatar: z.instanceof(File).optional(),
+	generateAvatar: z
+		.literal("on")
+		.optional()
+		.transform((el) => !!el),
 });
 
 const uploadAvatar = async (avatar: File, assistantId: ObjectId): Promise<string> => {
@@ -88,6 +95,29 @@ export const actions: Actions = {
 			}
 
 			hash = await uploadAvatar(parse.data.avatar, newAssistantId);
+		} else if (
+			ASSISTANTS_GENERATE_AVATAR === "true" &&
+			HF_TOKEN !== "" &&
+			parse.data.generateAvatar
+		) {
+			try {
+				const avatar = await timeout(
+					generateAvatar(parse.data.description, parse.data.name),
+					30000
+				);
+
+				hash = await uploadAvatar(avatar, newAssistantId);
+			} catch (err) {
+				return fail(400, {
+					error: true,
+					errors: [
+						{
+							field: "avatar",
+							message: "Avatar generation failed. Try again or disable the feature.",
+						},
+					],
+				});
+			}
 		}
 
 		const { insertedId } = await collections.assistants.insertOne({