Skip to content

Commit

Permalink
implement #23
Browse files Browse the repository at this point in the history
  • Loading branch information
ZeldaFan0225 committed Jan 27, 2023
1 parent 979e27c commit 3e8682c
Show file tree
Hide file tree
Showing 10 changed files with 276 additions and 6 deletions.
4 changes: 4 additions & 0 deletions changelog.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
# Changelog

## V2.3.3

- add remix context menu to remix users avatars with prompts

## V2.3.2

- add ability to gift kudos on /generate and /advanced_generate result messages
Expand Down
19 changes: 19 additions & 0 deletions config.md
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,25 @@ Here you can see an explanation of what which option does
"allow_sharing": (BOOLEAN) *5,
"allow_rating": Allow the user to rate their generated images (BOOLEAN)
}
},
"remix": {
"enabled": Whether the remix action is enabled (BOOLEAN),
"require_login": Set to true if you want the user to log in with their stable horde token (BOOLEAN),
"trusted_workers": Whether to only use trusted workers (BOOLEAN) *1,
"blacklisted_words": A list of blacklisted words which users are not allowed to use (ARRAY OF STRING),
"convert_a1111_weight_to_horde_weight": Whether to convert a1111 to weighted prompt required by the api (BOOLEAN) *7 *8,
"generation_options": {
"sampler_name": The sampler to use for remixing (STRING) *1 *2,
"model": The model to use for remixing (STRING) *1 *3,
"width": The width of the result (INTEGER) *1,
"height": The height of the result (INTEGER) *1,
"allow_nsfw": Set to true if you want to allow NSFW image generation for your users (BOOLEAN),
"censor_nsfw": Whether to censor NSFW images if they were generated by accident (BOOLEAN) *1,
"share_result": Whether to share the result or not (BOOLEAN) *1,
"cfg": The cfg scale to use for remixing (INTEGER) *1,
"denoise": The denoise strength to use for remixing (INTEGER; between 0 and 100) *1,
"steps": The steps to use for remixing (INTEGER)
}
}
}
```
Expand Down
2 changes: 1 addition & 1 deletion package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "zeldafan_discord_bot",
"version": "2.3.2",
"version": "2.3.3",
"description": "",
"main": "dist/index.js",
"scripts": {
Expand Down
5 changes: 3 additions & 2 deletions src/commands/advanced_generate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -509,8 +509,9 @@ ETA: <t:${Math.floor(Date.now()/1000)+(status?.wait_time ?? 0)}:R>`
const images = await ctx.stable_horde_manager.getGenerationStatus(generation_start!.id!)

const image_map_r = images.generations?.map(async (g, i) => {
const req = await Centra(g.img!, "get").send().then(res => res.body)
const attachment = new AttachmentBuilder(req, {name: `${g.seed ?? `image${i}`}.webp`})
const req = await Centra(g.img!, "get").send();
if(ctx.client.config.advanced?.dev) console.log(req)
const attachment = new AttachmentBuilder(req.body, {name: `${g.seed ?? `image${i}`}.webp`})
const embed = new EmbedBuilder({
title: `Image ${i+1}`,
image: {url: `attachment://${g.seed ?? `image${i}`}.webp`},
Expand Down
5 changes: 3 additions & 2 deletions src/commands/generate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -319,8 +319,9 @@ ETA: <t:${Math.floor(Date.now()/1000)+(status?.wait_time ?? 0)}:R>`
const images = await ctx.stable_horde_manager.getGenerationStatus(generation_start!.id!)

const image_map_r = images.generations?.map(async (g, i) => {
const req = await Centra(g.img!, "get").send().then(res => res.body)
const attachment = new AttachmentBuilder(req, {name: `${g.seed ?? `image${i}`}.webp`})
const req = await Centra(g.img!, "get").send();
if(ctx.client.config.advanced?.dev) console.log(req)
const attachment = new AttachmentBuilder(req.body, {name: `${g.seed ?? `image${i}`}.webp`})
const embed = new EmbedBuilder({
title: `Image ${i+1}`,
image: {url: `attachment://${g.seed ?? `image${i}`}.webp`},
Expand Down
57 changes: 57 additions & 0 deletions src/contexts/remix.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import { ApplicationCommandType, ContextMenuCommandBuilder } from "discord.js";
import { Context } from "../classes/context";
import { ContextContext } from "../classes/contextContext";


const command_data = new ContextMenuCommandBuilder()
.setType(ApplicationCommandType.User)
.setName("Remix Profile Picture")
.setDMPermission(false)

export default class extends Context {
constructor() {
super({
name: "Remix Profile Picture",
command_data: command_data.toJSON(),
staff_only: false,
})
}

override async run(ctx: ContextContext<ApplicationCommandType.User>): Promise<any> {
if(!ctx.database) return ctx.error({error: "The database is disabled. This action requires a database."})
if(!ctx.client.config.remix?.enabled) return ctx.error({error: "This feature has been disabled"})
const target_user = ctx.interaction.targetUser
const user_token = await ctx.client.getUserToken(ctx.interaction.user.id, ctx.database)

if(!user_token) return ctx.error({error: `You are required to ${await ctx.client.getSlashCommandTag("login")} to use this action`, codeblock: false})

const modal = {
title: "Enter Prompt",
custom_id: `remix_${target_user.id}`,
components: [{
type: 1,
components: [{
type: 4,
style: 1,
label: "Prompt",
custom_id: "prompt",
required: true,
placeholder: "Enter a prompt to remix the users avatar with"
}]
},{
type: 1,
components: [{
type: 4,
style: 1,
label: "Edit Strength",
custom_id: "strength",
required: false,
placeholder: "Number between 1 and 100",
max_length: 3
}]
}]
}

await ctx.interaction.showModal(modal)
}
}
150 changes: 150 additions & 0 deletions src/modals/remix.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
import Centra from "centra";
import { Modal } from "../classes/modal";
import { ModalContext } from "../classes/modalContext";
import { appendFileSync } from "fs";
import StableHorde from "@zeldafan0225/stable_horde";
import { EmbedBuilder } from "@discordjs/builders";
import { AttachmentBuilder, ButtonBuilder, Colors } from "discord.js";


export default class extends Modal {
constructor() {
super({
name: "remix",
staff_only: false,
regex: /remix_\d{17,20}/
})
}

override async run(ctx: ModalContext): Promise<any> {
if(!ctx.client.config.remix?.enabled) return ctx.error({error: "This feature has been disabled"})
if(!ctx.database) return ctx.error({error: "The database is disabled. This action requires a database."})

await ctx.interaction.deferReply()

const target_user_id = ctx.interaction.customId.split("_")[1] ?? ctx.interaction.user.id
const target_user = await ctx.client.users.fetch(target_user_id).catch(console.error)
const user_token = await ctx.client.getUserToken(ctx.interaction.user.id, ctx.database)
if(!target_user?.id) return ctx.error({error: "Unable to find target user"})
if(ctx.client.config.remix?.require_login && !user_token) return ctx.error({error: `You are required to ${await ctx.client.getSlashCommandTag("login")} to use the remix action`, codeblock: false})

const avatar_url = target_user.displayAvatarURL({size: 2048, extension: "webp"})
let prompt = ctx.interaction.fields.getTextInputValue("prompt")


const image_req = await Centra(avatar_url, "GET").send()
if(!image_req.statusCode?.toString().startsWith("2")) return ctx.error({error: "Unable to fetch users avatar"})
const image = image_req.body.toString("base64")


if(ctx.client.config.remix?.blacklisted_words?.some(w => prompt.toLowerCase().includes(w.toLowerCase()))) return ctx.error({error: "Your prompt included one or more blacklisted words"})

if(ctx.client.config.remix?.convert_a1111_weight_to_horde_weight) {
prompt = prompt.replace(/(\(+|\[+)|(\)+|\]+)|\:\d\.\d(\)+|\]+)/g, (w) => {
if(w.startsWith("(") || w.startsWith("[")) return "("
if(w.startsWith(":")) return w;
const weight = 1 + (0.1 * (w.startsWith(")") ? 1 : -1) * w.length)
return `:${weight.toFixed(1)})`
})
}

const generation_options = ctx.client.config.remix.generation_options

let strength = parseInt(ctx.interaction.fields.getTextInputValue("strength") || `${generation_options?.denoise ?? 100}` || "60")
if(strength > 100) strength = 100
if(strength < 1) strength = 1

const generation_data: StableHorde.GenerationInput = {
prompt,
params: {
sampler_name: generation_options?.sampler_name as typeof StableHorde.ModelGenerationInputStableSamplers[keyof typeof StableHorde.ModelGenerationInputStableSamplers] | undefined,
height: generation_options?.height,
width: generation_options?.width,
cfg_scale: generation_options?.cfg,
steps: generation_options?.steps,
denoising_strength: strength / 100
},
nsfw: !!(generation_options?.allow_nsfw ?? true),
censor_nsfw: !!(generation_options?.censor_nsfw ?? true),
trusted_workers: !!(ctx.client.config.remix?.trusted_workers ?? true),
models: generation_options?.model ? [generation_options.model] : undefined,
r2: true,
shared: generation_options?.share_result,
source_image: image,
source_processing: StableHorde.SourceImageProcessingTypes.img2img
}

if(ctx.client.config.advanced?.dev) {
console.log(generation_data)
console.log(generation_options)
}

const message = await ctx.interaction.editReply({content: "Remixing..."})

const generation_start = await ctx.stable_horde_manager.postAsyncGenerate(generation_data, {token: user_token})
.catch((e) => {
if(ctx.client.config.advanced?.dev) console.error(e)
return e;
})
if(!generation_start || !generation_start.id) return await error(generation_start.message)

if(ctx.client.config.advanced?.dev) {
console.log(generation_start)
}

if (ctx.client.config.logs?.enabled) {
if (ctx.client.config.logs.log_actions?.img2img) {
if (ctx.client.config.logs.plain) logGeneration("txt");
if (ctx.client.config.logs.csv) logGeneration("csv");
}
function logGeneration(type: "txt" | "csv") {
ctx.client.initLogDir();
const log_dir = ctx.client.config.logs?.directory ?? "/logs";
const content = type === "csv" ? `\n${new Date().toISOString()},${ctx.interaction.user.id},${generation_start?.id},${true},"${prompt}"` : `\n${new Date().toISOString()} | ${ctx.interaction.user.id}${" ".repeat(20 - ctx.interaction.user.id.length)} | ${generation_start?.id} | ${true}${" ".repeat(10)} | ${prompt}`;
appendFileSync(`${process.cwd()}${log_dir}/logs_${new Date().getMonth() + 1}-${new Date().getFullYear()}.${type}`, content);
}
}

const inter = setInterval(async () => {
const status = await ctx.stable_horde_manager.getGenerationCheck(generation_start?.id!)
if(ctx.client.config.advanced?.dev) console.log(status)
if(status.done) await displayResult()
else if(status.faulted) displayError()
}, 1000 * 5)

async function displayResult() {
clearInterval(inter)
if(!target_user?.id) return displayError();
const result = await ctx.stable_horde_manager.getGenerationStatus(generation_start.id)
const generation = result.generations?.[0]
if(!generation?.id) return displayError();
if(ctx.client.config.advanced?.dev) console.log(generation)
const req = await Centra(generation.img!, "get").send().then(res => res.body)
const attachment = new AttachmentBuilder(req, {name: `${generation.seed ?? `image${0}`}.webp`})
const embed = new EmbedBuilder({
color: Colors.Blue,
title: "Remixing finished",
description: `**Prompt**\n${prompt}\n**Target**\n${target_user?.tag}\n**Strength** ${strength}%${ctx.client.config.advanced?.dev ? `\n**Seed** ${generation.seed}` : ""}`,
thumbnail: {url: target_user?.displayAvatarURL({extension: "webp", size: 2048})},
image: {url: `attachment://${generation.seed ?? `image${0}`}.webp`},
})
embed.setThumbnail(`attachment://original.webp`)
const delete_btn = new ButtonBuilder({
label: "Delete this message",
custom_id: `delete_${ctx.interaction.user.id}`,
style: 4
})
await message.edit({content: null, files: [attachment], embeds: [embed], components: [{type: 1, components: [delete_btn.toJSON()]}]}).catch(console.error)
}

async function displayError() {
clearInterval(inter)
await error()
}

async function error(msg?: string) {
await ctx.interaction.followUp({content: `Unable to remix...${msg?.length ? `\n${msg}` : ""}`, ephemeral: true}).catch(console.error)
await message.delete()
}
}
}
19 changes: 19 additions & 0 deletions src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -254,5 +254,24 @@ export interface Config {
allow_sharing?: boolean,
allow_rating?: boolean
}
},
remix?: {
enabled?: boolean,
require_login?: boolean,
trusted_workers?: boolean,
blacklisted_words?: string[],
convert_a1111_weight_to_horde_weight?: boolean,
generation_options?: {
sampler_name?: string,
width?: number,
height?: number,
allow_nsfw?: boolean,
censor_nsfw?: boolean,
model: string,
share_result?: boolean,
cfg?: number,
denoise?: number,
steps?: number
}
}
}
19 changes: 19 additions & 0 deletions template.config.json
Original file line number Diff line number Diff line change
Expand Up @@ -174,5 +174,24 @@
"allow_sharing": true,
"allow_rating": true
}
},
"remix": {
"enabled": true,
"require_login": true,
"trusted_workers": true,
"blacklisted_words": [],
"convert_a1111_weight_to_horde_weight": true,
"generation_options": {
"sampler_name": "k_euler",
"model": "stable_diffusion",
"width": 512,
"height": 512,
"allow_nsfw": false,
"censor_nsfw": true,
"share_result": false,
"cfg": 7,
"denoise": 60,
"steps": 50
}
}
}

0 comments on commit 3e8682c

Please sign in to comment.