feat: send typing and uploading actions

README.md

@@ -22,34 +22,3 @@ You can put these in `.env` file or pass them as environment variables.
 
 - Start stable diffusion webui: `cd sd-webui`, `./webui.sh --api`
 - Start bot: `deno task start`
-
-## TODO
-
-- [x] Keep generation history
-- [x] Changing params, parsing png info in request
-- [x] Cancelling jobs by deleting message
-- [x] Multiple parallel workers
-- [x] Replying to another text message to copy prompt and generate
-- [x] Replying to bot message, conversation in DMs
-- [x] Replying to png message to extract png info nad generate
-- [ ] Banning tags
-- [x] Img2Img + Upscale
-  - [ ] special param "scale" to change image size preserving aspect ratio
-- [ ] Admin WebUI
-- [ ] User daily generation limits
-- [ ] Querying all generation history, displaying stats
-- [ ] Analyzing prompt quality based on tag csv
-- [ ] Report aliased/unknown tags based on csv
-- [ ] Report unknown loras
-- [ ] Investigate "sendMediaGroup failed"
-- [ ] Changing sampler without error on unknown sampler
-- [ ] Changing model
-- [ ] Inpaint using telegram photo edit
-- [ ] Outpaint
-- [ ] Non-SD (extras) upscale
-- [ ] Tiled generation to allow very big images
-- [ ] Downloading raw images
-- [ ] Extra prompt syntax, fixing `()+++` syntax
-- [ ] Translations
-  - replace fmtDuration usage
-  - replace formatOrdinal usage

bot/img2imgCommand.ts

@@ -115,8 +115,7 @@ async function img2img(
     from: ctx.message.from,
     chat: ctx.message.chat,
     requestMessageId: ctx.message.message_id,
-    replyMessageId: replyMessage.message_id,
-    status: { type: "waiting" },
+    status: { type: "waiting", message: replyMessage },
   });
 
   logger().debug(`Job enqueued for ${formatUserChat(ctx.message)}`);

bot/queueCommand.ts

@@ -16,6 +16,7 @@ export async function queueCommand(ctx: Grammy.CommandContext<Context>) {
       .then((jobs) => jobs.map((job, index) => ({ ...job.value, place: index + 1 })));
     const jobs = [...processingJobs, ...waitingJobs];
     const { bold } = GrammyParseMode;
+
     return fmt([
       "Current queue:\n",
       ...jobs.length > 0
@@ -47,8 +48,9 @@ export async function queueCommand(ctx: Grammy.CommandContext<Context>) {
   }
 
   async function handleFutureUpdates() {
-    for (let idx = 0; idx < 20; idx++) {
-      await new Promise((resolve) => setTimeout(resolve, 3000));
+    for (let idx = 0; idx < 30; idx++) {
+      await ctx.api.sendChatAction(ctx.chat.id, "typing");
+      await new Promise((resolve) => setTimeout(resolve, 4000));
       const nextFormattedMessage = await getMessageText();
       if (nextFormattedMessage.text !== formattedMessage.text) {
         await ctx.api.editMessageText(

bot/txt2imgCommand.ts

@@ -96,8 +96,7 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
     from: ctx.message.from,
     chat: ctx.message.chat,
     requestMessageId: ctx.message.message_id,
-    replyMessageId: replyMessage.message_id,
-    status: { type: "waiting" },
+    status: { type: "waiting", message: replyMessage },
   });
 
   logger().debug(`Job enqueued for ${formatUserChat(ctx.message)}`);

db/jobStore.ts

@@ -4,16 +4,36 @@ import { db } from "./db.ts";
 
 export interface JobSchema {
   task:
-    | { type: "txt2img"; params: Partial<PngInfo> }
-    | { type: "img2img"; params: Partial<PngInfo>; fileId: string };
+    | {
+      type: "txt2img";
+      params: Partial<PngInfo>;
+    }
+    | {
+      type: "img2img";
+      params: Partial<PngInfo>;
+      fileId: string;
+    };
   from: GrammyTypes.User;
   chat: GrammyTypes.Chat;
   requestMessageId: number;
-  replyMessageId?: number;
   status:
-    | { type: "waiting" }
-    | { type: "processing"; progress: number; worker: string; updatedDate: Date }
-    | { type: "done"; info?: SdTxt2ImgInfo; startDate?: Date; endDate?: Date };
+    | {
+      type: "waiting";
+      message?: GrammyTypes.Message.TextMessage;
+    }
+    | {
+      type: "processing";
+      progress: number;
+      worker: string;
+      updatedDate: Date;
+      message?: GrammyTypes.Message.TextMessage;
+    }
+    | {
+      type: "done";
+      info?: SdTxt2ImgInfo;
+      startDate?: Date;
+      endDate?: Date;
+    };
 }
 
 export const jobStore = new IKV.Store(db, "job", {

tasks/processJobs.ts

@@ -40,10 +40,16 @@ export async function processJobs(): Promise<never> {
       if (!worker) continue;
 
       // process the job
-      await job.update({
-        status: { type: "processing", progress: 0, worker: worker.name, updatedDate: new Date() },
-      });
-
+      await job.update((value) => ({
+        ...value,
+        status: {
+          type: "processing",
+          progress: 0,
+          worker: worker.name,
+          updatedDate: new Date(),
+          message: job.value.status.type !== "done" ? job.value.status.message : undefined,
+        },
+      }));
       busyWorkers.add(worker.name);
       processJob(job, worker, config)
         .catch(async (err) => {
@@ -89,31 +95,63 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
   );
   const startDate = new Date();
 
-  // if there is already a status message delete it
-  if (job.value.replyMessageId) {
-    await bot.api.deleteMessage(job.value.chat.id, job.value.replyMessageId)
-      .catch(() => undefined);
+  // if there is already a status message and its older than 10 seconds
+  if (
+    job.value.status.type === "processing" && job.value.status.message &&
+    (Date.now() - job.value.status.message.date * 1000) > 10 * 1000
+  ) {
+    // delete it
+    await bot.api.deleteMessage(
+      job.value.status.message.chat.id,
+      job.value.status.message.message_id,
+    ).catch(() => undefined);
+    await job.update((value) => ({
+      ...value,
+      status: { ...value.status, message: undefined },
+    }));
   }
 
-  // send a new status message
-  const newStatusMessage = await bot.api.sendMessage(
-    job.value.chat.id,
-    `Generating your prompt now... 0% using ${worker.name}`,
-    { reply_to_message_id: job.value.requestMessageId },
-  ).catch((err) => {
-    // don't error if the request message was deleted
-    if (err instanceof Grammy.GrammyError && err.message.match(/repl(y|ied)/)) return null;
-    else throw err;
-  });
-  // if the request message was deleted, cancel the job
-  if (!newStatusMessage) {
+  // we have to check if job is still processing at every step because TypeScript
+  if (job.value.status.type === "processing") {
+    await bot.api.sendChatAction(job.value.chat.id, "upload_photo", { maxAttempts: 1 })
+      .catch(() => undefined);
+    // if now there is no status message
+    if (!job.value.status.message) {
+      // send a new status message
+      const statusMessage = await bot.api.sendMessage(
+        job.value.chat.id,
+        `Generating your prompt now... 0% using ${worker.name}`,
+        { reply_to_message_id: job.value.requestMessageId },
+      ).catch((err) => {
+        // if the request message (the message we are replying to) was deleted
+        if (err instanceof Grammy.GrammyError && err.message.match(/repl(y|ied)/)) {
+          // jest set the status message to undefined
+          return undefined;
+        }
+        throw err;
+      });
+      await job.update((value) => ({
+        ...value,
+        status: { ...value.status, message: statusMessage },
+      }));
+    } else {
+      // edit the existing status message
+      await bot.api.editMessageText(
+        job.value.status.message.chat.id,
+        job.value.status.message.message_id,
+        `Generating your prompt now... 0% using ${worker.name}`,
+        { maxAttempts: 1 },
+      ).catch(() => undefined);
+    }
+  }
+
+  // if we don't have a status message (it failed sending because request was deleted)
+  if (job.value.status.type === "processing" && !job.value.status.message) {
+    // cancel the job
     await job.delete();
-    logger().info(
-      `Job cancelled for ${formatUserChat(job.value)}`,
-    );
+    logger().info(`Job cancelled for ${formatUserChat(job.value)}`);
     return;
   }
-  await job.update({ replyMessageId: newStatusMessage.message_id });
 
   // reduce size if worker can't handle the resolution
   const size = limitSize(
@@ -123,25 +161,29 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
 
   // process the job
   const handleProgress = async (progress: SdProgressResponse) => {
-    // important: don't let any errors escape this callback
-    if (job.value.replyMessageId) {
-      await bot.api.editMessageText(
-        job.value.chat.id,
-        job.value.replyMessageId,
-        `Generating your prompt now... ${
-          (progress.progress * 100).toFixed(0)
-        }% using ${worker.name}`,
-        { maxAttempts: 1 },
-      ).catch(() => undefined);
+    if (job.value.status.type === "processing" && job.value.status.message) {
+      await Promise.all([
+        bot.api.sendChatAction(job.value.chat.id, "upload_photo", { maxAttempts: 1 }),
+        bot.api.editMessageText(
+          job.value.status.message.chat.id,
+          job.value.status.message.message_id,
+          `Generating your prompt now... ${
+            (progress.progress * 100).toFixed(0)
+          }% using ${worker.name}`,
+          { maxAttempts: 1 },
+        ),
+        job.update((value) => ({
+          ...value,
+          status: {
+            type: "processing",
+            progress: progress.progress,
+            worker: worker.name,
+            updatedDate: new Date(),
+            message: value.status.type !== "done" ? value.status.message : undefined,
+          },
+        }), { maxAttempts: 1 }),
+      ]).catch(() => undefined);
     }
-    await job.update({
-      status: {
-        type: "processing",
-        progress: progress.progress,
-        worker: worker.name,
-        updatedDate: new Date(),
-      },
-    }, { maxAttempts: 1 }).catch(() => undefined);
   };
   let response: SdResponse<unknown>;
   const taskType = job.value.task.type; // don't narrow this to never pls typescript
@@ -169,11 +211,11 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
       throw new Error(`Unknown task type: ${taskType}`);
   }
 
-  // upload the result
-  if (job.value.replyMessageId) {
+  // change status message to uploading images
+  if (job.value.status.type === "processing" && job.value.status.message) {
     await bot.api.editMessageText(
-      job.value.chat.id,
-      job.value.replyMessageId,
+      job.value.status.message.chat.id,
+      job.value.status.message.message_id,
       `Uploading your images...`,
     ).catch(() => undefined);
   }
@@ -201,10 +243,13 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
       : [],
   ]);
 
+  // sending images loop because telegram is unreliable and it would be a shame to lose the images
   let sendMediaAttempt = 0;
   let resultMessages: GrammyTypes.Message.MediaMessage[] | undefined;
   while (true) {
     sendMediaAttempt++;
+    await bot.api.sendChatAction(job.value.chat.id, "upload_photo", { maxAttempts: 1 })
+      .catch(() => undefined);
 
     // parse files from reply JSON
     const inputFiles = await Promise.all(
@@ -232,7 +277,12 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
     } catch (err) {
       logger().warning(`Sending images (attempt ${sendMediaAttempt}) failed: ${err}`);
       if (sendMediaAttempt >= 6) throw err;
-      await Async.delay(10000);
+      // wait 2 * 5 seconds before retrying
+      for (let i = 0; i < 2; i++) {
+        await bot.api.sendChatAction(job.value.chat.id, "upload_photo", { maxAttempts: 1 })
+          .catch(() => undefined);
+        await Async.delay(5000);
+      }
     }
   }
 
@@ -245,17 +295,23 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
   }
 
   // delete the status message
-  if (job.value.replyMessageId) {
-    await bot.api.deleteMessage(job.value.chat.id, job.value.replyMessageId)
-      .catch(() => undefined)
-      .then(() => job.update({ replyMessageId: undefined }))
-      .catch(() => undefined);
+  if (job.value.status.type === "processing" && job.value.status.message) {
+    await bot.api.deleteMessage(
+      job.value.status.message.chat.id,
+      job.value.status.message.message_id,
+    ).catch(() => undefined);
+    await job.update((value) => ({
+      ...value,
+      status: { ...value.status, message: undefined },
+    }));
   }
 
   // update job to status done
-  await job.update({
+  await job.update((value) => ({
+    ...value,
     status: { type: "done", info: response.info, startDate, endDate: new Date() },
-  });
+  }));
+
   logger().debug(
     `Job finished for ${formatUserChat(job.value)} using ${worker.name}${
       sendMediaAttempt > 1 ? ` after ${sendMediaAttempt} attempts` : ""

tasks/updateJobStatusMsgs.ts

@@ -14,10 +14,10 @@ export async function updateJobStatusMsgs(): Promise<never> {
       await new Promise((resolve) => setTimeout(resolve, 5000));
       const jobs = await jobStore.getBy("status.type", "waiting");
       for (const [index, job] of jobs.entries()) {
-        if (!job.value.replyMessageId) continue;
+        if (job.value.status.type !== "waiting" || !job.value.status.message) continue;
         await bot.api.editMessageText(
-          job.value.chat.id,
-          job.value.replyMessageId,
+          job.value.status.message.chat.id,
+          job.value.status.message.message_id,
           `You are ${formatOrdinal(index + 1)} in queue.`,
           { maxAttempts: 1 },
         ).catch(() => undefined);