feat: add scale parameter to img2img

This commit is contained in:
pinks 2023-09-16 13:49:12 +02:00
parent 8e81f82d8b
commit 0517ce1930
4 changed files with 41 additions and 40 deletions

View File

@ -73,22 +73,10 @@ async function img2img(
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) {
// TODO: remove bot command from replied to text
const originalParams = parsePngInfo(repliedToText);
params = {
...originalParams,
...params,
prompt: [originalParams.prompt, params.prompt].filter(Boolean).join("\n"),
negative_prompt: [originalParams.negative_prompt, params.negative_prompt]
.filter(Boolean).join("\n"),
};
params = parsePngInfo(repliedToText, params);
}
const messageParams = parsePngInfo(match ?? "");
params = {
...params,
...messageParams,
prompt: [params.prompt, messageParams.prompt].filter(Boolean).join("\n"),
};
params = parsePngInfo(match ?? "", params);
if (!fileId) {
await ctx.reply(

View File

@ -50,35 +50,16 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
if (includeRepliedTo && repliedToMsg?.document?.mime_type === "image/png") {
const file = await ctx.api.getFile(repliedToMsg.document.file_id);
const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer());
const fileParams = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "");
params = {
...params,
...fileParams,
prompt: [params.prompt, fileParams.prompt].filter(Boolean).join("\n"),
negative_prompt: [params.negative_prompt, fileParams.negative_prompt]
.filter(Boolean).join("\n"),
};
params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "", params);
}
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) {
// TODO: remove bot command from replied to text
const originalParams = parsePngInfo(repliedToText);
params = {
...originalParams,
...params,
prompt: [originalParams.prompt, params.prompt].filter(Boolean).join("\n"),
negative_prompt: [originalParams.negative_prompt, params.negative_prompt]
.filter(Boolean).join("\n"),
};
params = parsePngInfo(repliedToText, params);
}
const messageParams = parsePngInfo(match);
params = {
...params,
...messageParams,
prompt: [params.prompt, messageParams.prompt].filter(Boolean).join("\n"),
};
params = parsePngInfo(match, params);
if (!params.prompt) {
await ctx.reply(

View File

@ -20,10 +20,14 @@ export interface PngInfo {
denoising_strength: number;
}
export function parsePngInfo(pngInfo: string): Partial<PngInfo> {
interface PngInfoExtra extends PngInfo {
upscale?: number;
}
export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>): Partial<PngInfo> {
const tags = pngInfo.split(/[,;]+|\.+\s|\n/u);
let part: "prompt" | "negative_prompt" | "params" = "prompt";
const params: Partial<PngInfo> = {};
const params: Partial<PngInfoExtra> = {};
const prompt: string[] = [];
const negativePrompt: string[] = [];
for (const tag of tags) {
@ -72,6 +76,13 @@ export function parsePngInfo(pngInfo: string): Partial<PngInfo> {
}
break;
}
case "upscale":
case "scale": {
part = "params";
const upscale = Number(value.trim());
if (upscale > 0) params.upscale = Math.min(upscale, 2);
break;
}
case "denoisingstrength":
case "denoising":
case "denoise": {
@ -113,5 +124,19 @@ export function parsePngInfo(pngInfo: string): Partial<PngInfo> {
}
if (prompt.length > 0) params.prompt = prompt.join(", ");
if (negativePrompt.length > 0) params.negative_prompt = negativePrompt.join(", ");
return params;
// handle upscale
if (params.upscale && baseParams?.width && baseParams?.height) {
params.width = baseParams.width * params.upscale;
params.height = baseParams.height * params.upscale;
}
return {
...baseParams,
...params,
prompt: [baseParams?.prompt, params.prompt]
.filter(Boolean).join("\n"),
negative_prompt: [baseParams?.negative_prompt, params.negative_prompt]
.filter(Boolean).join("\n"),
};
}

View File

@ -213,7 +213,14 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
case "txt2img":
response = await sdTxt2Img(
worker.api,
{ ...config.defaultParams, ...job.value.task.params, ...size },
{
...config.defaultParams,
...job.value.task.params,
...size,
negative_prompt: job.value.task.params.negative_prompt
? job.value.task.params.negative_prompt
: config.defaultParams?.negative_prompt,
},
handleProgress,
);
break;