README.md

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

bot/queueCommand.ts

@@ -16,7 +16,6 @@ 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
@@ -48,9 +47,8 @@ export async function queueCommand(ctx: Grammy.CommandContext<Context>) {
   }
 
   async function handleFutureUpdates() {
-    for (let idx = 0; idx < 30; idx++) {
-      await ctx.api.sendChatAction(ctx.chat.id, "typing");
-      await new Promise((resolve) => setTimeout(resolve, 4000));
+    for (let idx = 0; idx < 20; idx++) {
+      await new Promise((resolve) => setTimeout(resolve, 3000));
       const nextFormattedMessage = await getMessageText();
       if (nextFormattedMessage.text !== formattedMessage.text) {
         await ctx.api.editMessageText(

bot/txt2imgCommand.ts

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

db/jobStore.ts

@@ -4,36 +4,16 @@ 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";
-      message?: GrammyTypes.Message.TextMessage;
-    }
-    | {
-      type: "processing";
-      progress: number;
-      worker: string;
-      updatedDate: Date;
-      message?: GrammyTypes.Message.TextMessage;
-    }
-    | {
-      type: "done";
-      info?: SdTxt2ImgInfo;
-      startDate?: Date;
-      endDate?: Date;
-    };
+    | { type: "waiting" }
+    | { type: "processing"; progress: number; worker: string; updatedDate: Date }
+    | { type: "done"; info?: SdTxt2ImgInfo; startDate?: Date; endDate?: Date };
 }
 
 export const jobStore = new IKV.Store(db, "job", {

tasks/processJobs.ts

@@ -40,16 +40,10 @@ export async function processJobs(): Promise<never> {
       if (!worker) continue;
 
       // process the job
-      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,
-        },
-      }));
+      await job.update({
+        status: { type: "processing", progress: 0, worker: worker.name, updatedDate: new Date() },
+      });
+
       busyWorkers.add(worker.name);
       processJob(job, worker, config)
         .catch(async (err) => {
@@ -95,63 +89,31 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
   );
   const startDate = new Date();
 
-  // 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 },
-    }));
+  // 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);
   }
 
-  // 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(
+  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) => {
-        // 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;
+    // 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;
   });
-      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
+  // if the request message was deleted, cancel the job
+  if (!newStatusMessage) {
     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(
@@ -161,29 +123,25 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
 
   // process the job
   const handleProgress = async (progress: SdProgressResponse) => {
-    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,
+    // 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 },
-        ),
-        job.update((value) => ({
-          ...value,
+      ).catch(() => undefined);
+    }
+    await job.update({
       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);
-    }
+    }, { maxAttempts: 1 }).catch(() => undefined);
   };
   let response: SdResponse<unknown>;
   const taskType = job.value.task.type; // don't narrow this to never pls typescript
@@ -211,11 +169,11 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
       throw new Error(`Unknown task type: ${taskType}`);
   }
 
-  // change status message to uploading images
-  if (job.value.status.type === "processing" && job.value.status.message) {
+  // upload the result
+  if (job.value.replyMessageId) {
     await bot.api.editMessageText(
-      job.value.status.message.chat.id,
-      job.value.status.message.message_id,
+      job.value.chat.id,
+      job.value.replyMessageId,
       `Uploading your images...`,
     ).catch(() => undefined);
   }
@@ -243,13 +201,10 @@ 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(
@@ -277,12 +232,7 @@ 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;
-      // 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);
-      }
+      await Async.delay(10000);
     }
   }
 
@@ -295,23 +245,17 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
   }
 
   // delete the status message
-  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 },
-    }));
+  if (job.value.replyMessageId) {
+    await bot.api.deleteMessage(job.value.chat.id, job.value.replyMessageId)
+      .catch(() => undefined)
+      .then(() => job.update({ replyMessageId: undefined }))
+      .catch(() => undefined);
   }
 
   // update job to status done
-  await job.update((value) => ({
-    ...value,
+  await job.update({
     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.status.type !== "waiting" || !job.value.status.message) continue;
+        if (!job.value.replyMessageId) continue;
         await bot.api.editMessageText(
-          job.value.status.message.chat.id,
-          job.value.status.message.message_id,
+          job.value.chat.id,
+          job.value.replyMessageId,
           `You are ${formatOrdinal(index + 1)} in queue.`,
           { maxAttempts: 1 },
         ).catch(() => undefined);