Compare commits

...

3 Commits

Author SHA1 Message Date
pinks 7541deb178 feat: send typing and uploading actions 2023-09-12 22:39:53 +02:00
pinks 80bcd539bd remove TODOs from readme 2023-09-12 22:37:55 +02:00
pinks c3e5a4d908 fix: reuse msgs instead of sending and deleting 2023-09-12 22:37:41 +02:00
7 changed files with 145 additions and 100 deletions

View File

@ -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

View File

@ -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)}`);

View File

@ -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(

View File

@ -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)}`);

View File

@ -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", {

View File

@ -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` : ""

View File

@ -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);