55 lines
1.7 KiB
TypeScript
55 lines
1.7 KiB
TypeScript
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 }; |