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; const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) { if (includeRepliedTo && repliedToText) {
// TODO: remove bot command from replied to text // TODO: remove bot command from replied to text
const originalParams = parsePngInfo(repliedToText); params = parsePngInfo(repliedToText, params);
params = {
...originalParams,
...params,
prompt: [originalParams.prompt, params.prompt].filter(Boolean).join("\n"),
negative_prompt: [originalParams.negative_prompt, params.negative_prompt]
.filter(Boolean).join("\n"),
};
} }
const messageParams = parsePngInfo(match ?? ""); params = parsePngInfo(match ?? "", params);
params = {
...params,
...messageParams,
prompt: [params.prompt, messageParams.prompt].filter(Boolean).join("\n"),
};
if (!fileId) { if (!fileId) {
await ctx.reply( 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") { if (includeRepliedTo && repliedToMsg?.document?.mime_type === "image/png") {
const file = await ctx.api.getFile(repliedToMsg.document.file_id); const file = await ctx.api.getFile(repliedToMsg.document.file_id);
const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer()); const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer());
const fileParams = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? ""); params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "", params);
params = {
...params,
...fileParams,
prompt: [params.prompt, fileParams.prompt].filter(Boolean).join("\n"),
negative_prompt: [params.negative_prompt, fileParams.negative_prompt]
.filter(Boolean).join("\n"),
};
} }
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption; const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) { if (includeRepliedTo && repliedToText) {
// TODO: remove bot command from replied to text // TODO: remove bot command from replied to text
const originalParams = parsePngInfo(repliedToText); params = parsePngInfo(repliedToText, params);
params = {
...originalParams,
...params,
prompt: [originalParams.prompt, params.prompt].filter(Boolean).join("\n"),
negative_prompt: [originalParams.negative_prompt, params.negative_prompt]
.filter(Boolean).join("\n"),
};
} }
const messageParams = parsePngInfo(match); params = parsePngInfo(match, params);
params = {
...params,
...messageParams,
prompt: [params.prompt, messageParams.prompt].filter(Boolean).join("\n"),
};
if (!params.prompt) { if (!params.prompt) {
await ctx.reply( await ctx.reply(

View File

@ -20,10 +20,14 @@ export interface PngInfo {
denoising_strength: number; 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); const tags = pngInfo.split(/[,;]+|\.+\s|\n/u);
let part: "prompt" | "negative_prompt" | "params" = "prompt"; let part: "prompt" | "negative_prompt" | "params" = "prompt";
const params: Partial<PngInfo> = {}; const params: Partial<PngInfoExtra> = {};
const prompt: string[] = []; const prompt: string[] = [];
const negativePrompt: string[] = []; const negativePrompt: string[] = [];
for (const tag of tags) { for (const tag of tags) {
@ -72,6 +76,13 @@ export function parsePngInfo(pngInfo: string): Partial<PngInfo> {
} }
break; 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 "denoisingstrength":
case "denoising": case "denoising":
case "denoise": { case "denoise": {
@ -113,5 +124,19 @@ export function parsePngInfo(pngInfo: string): Partial<PngInfo> {
} }
if (prompt.length > 0) params.prompt = prompt.join(", "); if (prompt.length > 0) params.prompt = prompt.join(", ");
if (negativePrompt.length > 0) params.negative_prompt = negativePrompt.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": case "txt2img":
response = await sdTxt2Img( response = await sdTxt2Img(
worker.api, 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, handleProgress,
); );
break; break;