Initial commit

src/handlers/imagine/imagine.ts

@@ -0,0 +1,55 @@
+import { Context } from "#root/context.js";
+import { Composer, InputFile, InputMediaBuilder } from "grammy";
+import { predict, sdClient } from "./sd-client.js";
+
+const composer = new Composer<Context>();
+const feature = composer.chatType("private");
+
+type StatusMessage = Awaited<ReturnType<Context["replyWithPhoto"]>>;
+
+async function sendProgress(ctx: Context) {
+    let isStatusSent = false;
+    let statusMessage: StatusMessage;
+
+    const progressInterval = setInterval(async () => {
+        const response = await sdClient.getProgress();
+
+        if (response.progress === 0.0 && response.state.job_count === 0) {
+            await statusMessage.delete();
+            clearInterval(progressInterval);
+        }
+
+        let buffer;
+
+        try {
+            buffer = Buffer.from(response.current_image, 'base64');
+        } catch (err) {
+            return;
+        }
+
+        if (!isStatusSent) {
+            statusMessage = await ctx.replyWithPhoto(new InputFile(buffer), {
+                caption: "Processing your image, please wait"
+            });
+
+            isStatusSent = true;
+            return;
+        } else {
+            const newMedia = InputMediaBuilder.photo(new InputFile(buffer), { caption: "Processing your image, please wait" });
+            await statusMessage.editMedia(newMedia);
+            return;
+        }
+
+    }, 700);
+}
+
+feature.command("imagine", async (ctx) => {
+    const prompt = ctx.match;
+    console.log(prompt);
+    await sendProgress(ctx);
+    const prediction = await predict(`mj ${prompt}`);
+    const image = await prediction.image.jpeg().toBuffer();
+    await ctx.replyWithPhoto(new InputFile(image));
+});
+
+export { composer as imageFeature };