Add base bot setup

Browse Source 
- supports ping, llama and image/comfyUI
 - comes with 3 pre-setup comfyUI templates
This commit is contained in:
BordedDev 2025-03-06 00:36:52 +01:00
commit 1f01c0c454
No known key found for this signature in database
GPG Key ID: C5F495EAE56673BF
16 changed files with 3114 additions and 0 deletions

132
.gitignore vendored Normal file
View File

@ -0,0 +1,132 @@
/.idea/
/.vscode/
/node_modules
.env
*.orig
*.pyc
*.swp
*~
.fuse_hidden*
.directory
.Trash-*
.nfs*
.idea/**/workspace.xml
.idea/**/tasks.xml
.idea/**/usage.statistics.xml
.idea/**/dictionaries
.idea/**/shelf
.idea/**/aws.xml
.idea/**/contentModel.xml
.idea/**/dataSources/
.idea/**/dataSources.ids
.idea/**/dataSources.local.xml
.idea/**/sqlDataSources.xml
.idea/**/dynamic.xml
.idea/**/uiDesigner.xml
.idea/**/dbnavigator.xml
.idea/**/gradle.xml
.idea/**/libraries
cmake-build-*/
.idea/**/mongoSettings.xml
*.iws
out/
.idea_modules/
atlassian-ide-plugin.xml
.idea/replstate.xml
.idea/sonarlint/
com_crashlytics_export_strings.xml
crashlytics.properties
crashlytics-build.properties
fabric.properties
.idea/httpRequests
.idea/caches/build_file_checksums.ser
logs
*.log
npm-debug.log*
yarn-debug.log*
yarn-error.log*
lerna-debug.log*
.pnpm-debug.log*
report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json
pids
*.pid
*.seed
*.pid.lock
lib-cov
coverage
*.lcov
.nyc_output
.grunt
bower_components
.lock-wscript
build/Release
node_modules/
jspm_packages/
web_modules/
*.tsbuildinfo
.npm
.eslintcache
.stylelintcache
.rpt2_cache/
.rts2_cache_cjs/
.rts2_cache_es/
.rts2_cache_umd/
.node_repl_history
*.tgz
.yarn-integrity
.env.development.local
.env.test.local
.env.production.local
.env.local
.cache
.parcel-cache
.next
out
.nuxt
dist
.cache/
.vuepress/dist
.temp
.docusaurus
.serverless/
.fusebox/
.dynamodb/
.tern-port
.vscode-test
.yarn/cache
.yarn/unplugged
.yarn/build-state.yml
.yarn/install-state.gz
.pnp.*
.svelte-kit/
package
Thumbs.db
Thumbs.db:encryptable
ehthumbs.db
ehthumbs_vista.db
*.stackdump
[Dd]esktop.ini
$RECYCLE.BIN/
*.cab
*.msi
*.msix
*.msm
*.msp
*.lnk
.DS_Store
.AppleDouble
.LSOverride
Icon
._*
.DocumentRevisions-V100
.fseventsd
.Spotlight-V100
.TemporaryItems
.Trashes
.VolumeIcon.icns
.com.apple.timemachine.donotpresent
.AppleDB
.AppleDesktop
Network Trash Folder
Temporary Items
.apdisk

24
deno.json Normal file
View File

@ -0,0 +1,24 @@
{
"tasks": {
"dev:snek-llama-bot": "deno run --allow-env --allow-read --allow-sys --allow-ffi --allow-net=snek.molodetz.nl:443,127.0.0.1:8188 --allow-run=\"deno\" src/ws-snek-llama-bot.ts",
"dev:snek-img-bot": "deno run --allow-env --allow-read --allow-sys --allow-ffi --allow-net=snek.molodetz.nl:443,127.0.0.1:8188 --allow-run=\"deno\" src/ws-snek-image-bot.ts"
},
"imports": {
"@eta-dev/eta": "jsr:@eta-dev/eta@^3.5.0",
"@logtape/logtape": "jsr:@logtape/logtape@^0.8.1",
"@std/assert": "jsr:@std/assert@1",
"@std/cli": "jsr:@std/cli@^1.0.12",
"@std/collections": "jsr:@std/collections@^1.0.10",
"@std/dotenv": "jsr:@std/dotenv@^0.225.3",
"@std/path": "jsr:@std/path@^1.0.8",
"@types/node": "npm:@types/node@^22.13.9",
"jmespath": "npm:jmespath@^0.16.0",
"lodash-es": "npm:lodash-es@^4.17.21",
"node-llama-cpp": "npm:node-llama-cpp@^3.6.0",
"octokit": "npm:octokit@^4.1.0"
},
"fmt": {
"semiColons": false
},
"nodeModulesDir": "auto"
}

1259
deno.lock Normal file
View File

File diff suppressed because it is too large Load Diff

41
src/msg-handlers/base-handler.ts Normal file
View File

@ -0,0 +1,41 @@
import { EventEmitter } from "node:events"
import {Bot, Message, User} from "../snek/snek-socket.ts"
import { trim } from "npm:lodash-es"
export abstract class BaseHandler extends EventEmitter {
prefix = "msg-handler"
user: User|null = null
constructor(prefix: string) {
super()
this.prefix = prefix.toLowerCase()
}
async bind(bot: Bot) {
this.user = await bot.user
bot.on("message", async (message) => {
try {
if (await this.isMatch(message)) {
try {
await this.handleMessage(message, bot)
} catch (e) {
message.reply("An error occurred while handling your message")
console.error("Error handling message", e)
}
}
} catch (e) {
console.error("Error checking message", e)
}
})
}
async isMatch(message: Message): Promise<boolean> {
return message.userUID !== this.user?.uid &&
trim(message?.message, ' `').toLowerCase().startsWith(this.prefix)
}
abstract handleMessage(
message: Message,
bot: Bot,
): Promise<void> | void
}

187
src/msg-handlers/img-gen-handler.ts Normal file
View File

@ -0,0 +1,187 @@
import {Bot, Message} from "../snek/snek-socket.ts"
import {BaseHandler} from "./base-handler.ts"
import {ImgGen} from "../util/img-gen.ts"
import {Eta} from "jsr:@eta-dev/eta"
import { getLogger } from "@logtape/logtape"
import * as path from "node:path"
import {randomUUID} from "node:crypto";
const TEMPLATES = new Eta({ views: path.join(Deno.cwd(), "templates/img-gen") })
const FORMATTING_WRAPPER = /^\s*`(?:``\w+)?(.*?)(?:``)?`\s*$/gs
const logger = getLogger(["img-gen-handler"])
const parsePrompt = (prompt: string): Record<string, unknown> => {
try {
return JSON.parse(prompt)
} catch (e) {
prompt = prompt.replace(FORMATTING_WRAPPER, "$1")
return JSON.parse(prompt)
}
}
export class ImgGenHandler extends BaseHandler {
#imageGenerator = new ImgGen()
#activeTemplate = "flux-dev.eta"
#templateVariables: Record<string, unknown> = {
steps: 20,
seed: 0,
cfg: 2,
sampler: "euler",
negativePrompt: "",
batchSize: 1,
width: 1024,
height: 1024,
}
#subCommands = {
"prompt": this.prompt.bind(this),
"template": this.template.bind(this),
"variables": this.variables.bind(this),
} as Record<string, (command: string, message: Message, bot: Bot) => void>
constructor() {
super("/img-gen")
}
override async handleMessage(message: Message, bot: Bot) {
const newMessage = message.message.substring(this.prefix.length).trim()
if (!newMessage) {
message.reply(
"No command given, try one of: " +
Object.keys(this.#subCommands).join(", "),
)
return
}
const [[_, command, rest]] = newMessage.matchAll(/^(\w+)\s*(.*)$/gs)
if (command in this.#subCommands) {
this.#subCommands[command]?.(rest.trim(), message, bot)
return
}
message.reply(
"Invalid command, assuming prompt. Otherwise these are the correct ones: " +
Object.keys(this.#subCommands).join(", "),
)
await this.prompt(newMessage, message, bot)
}
async prompt(command: string, message: Message, bot: Bot) {
if (!command) {
message.reply("No prompt given")
return
}
try {
logger.info("Generating image", { command, variables: this.#templateVariables, message })
const prompt = parsePrompt(
await TEMPLATES.renderAsync(this.#activeTemplate, {
...this.#templateVariables,
randomSeed: (
min: number = 0,
max: number = Number.MAX_SAFE_INTEGER,
) => Math.floor(Math.random() * (max - min + 1)) + min,
prompt: command.replaceAll("\n", " "),
}),
)
message.reply("image generated called")
const promptResults = await this.#imageGenerator.dispatchPrompt(
prompt,
)
const blob = promptResults.filter((result) => result instanceof Blob)
if (blob.length === 0) {
console.log("Prompt Results: ", promptResults)
message.reply("Failed to generate image")
return
}
const files = await Promise.all(
blob.map(async (blob) =>
new File([(await blob.arrayBuffer()).slice(8)], `${randomUUID()}.png`)
),
)
await bot.uploadFiles(message.channelUID, ...files)
} catch (e) {
console.error(e)
message.reply(`Failed to generate image: ${e.message}`)
}
}
async template(command: string, message: Message, bot: Bot) {
if (!command) {
const templatePath = TEMPLATES.config.views
if (templatePath) {
const templateOptions = await Deno.readDir(templatePath)
let templateList =
`Current template: ${this.#activeTemplate}\n\nTemplates:\n`
for await (const template of templateOptions) {
templateList += ` - ${template.name}\n`
}
message.reply(templateList)
} else {
message.reply(`Current template: ${this.#activeTemplate}`)
return
}
return
}
const [[_, template, rest]] = command.matchAll(/^(\S+)\s*(.*)$/gs)
if (!template) {
message.reply("No template given")
return
}
if (template.startsWith("@")) {
TEMPLATES.loadTemplate(template, rest, { async: true })
this.#activeTemplate = template
message.reply(`Template set to ${template}`)
} else {
try {
Deno.readFileSync(TEMPLATES.resolvePath(template))
this.#activeTemplate = template
message.reply(`Template set to ${template}`)
} catch (e) {
message.reply(`Failed to load template: ${e.message}`)
return
}
}
}
variables(command: string, message: Message, bot: Bot) {
if (!command) {
let currentVariables = "Current Variables:\n\n```json\n"
currentVariables += JSON.stringify(this.#templateVariables, null, 2)
currentVariables += "\n```"
message.reply(currentVariables)
return
}
const [[_, variable, value]] = command.matchAll(/^(\S+)\s*(.*)$/gs)
if (!variable) {
message.reply("No variable given")
return
}
if (!value) {
message.reply(
`Variable ${variable} = ${this.#templateVariables[variable]}`,
)
return
}
this.#templateVariables[variable] = value
message.reply(`Variable set: ${variable} = ${value}`)
}
}

212
src/msg-handlers/llama-handler.ts Normal file
View File

@ -0,0 +1,212 @@
import {BaseHandler} from "./base-handler.ts"
import {Bot, Message} from "../snek/snek-socket.ts"
import {trim, trimStart} from "npm:lodash-es"
import {
ChatSessionModelFunctions,
ChatWrapper,
GeneralChatWrapper,
getLlama,
LLamaChatPromptOptions,
LlamaChatSession,
LlamaContext,
LlamaModel,
resolveChatWrapper,
Token,
} from "npm:node-llama-cpp"
import {getLogger} from "@logtape/logtape"
import {deepMerge} from "@std/collections/deep-merge"
const llama = await getLlama()
const textEncoder = new TextEncoder()
function printSync(input: string | Uint8Array, to = Deno.stdout) {
let bytesWritten = 0
const bytes = typeof input === "string" ? textEncoder.encode(input) : input
while (bytesWritten < bytes.length) {
bytesWritten += to.writeSync(bytes.subarray(bytesWritten))
}
}
const logger = getLogger(["llama-gen-handler"])
const optionsGenerator = <
const Functions extends ChatSessionModelFunctions | undefined = undefined,
LLamaOptions = LLamaChatPromptOptions<Functions>
>(
model: LlamaModel,
debugOutput: boolean = true,
defaultTimeout = 5 * 60 * 1000,
options?: LLamaOptions,
): LLamaChatPromptOptions<Functions> => {
const manager = AbortSignal.timeout(defaultTimeout)
const defaultOptions: LLamaChatPromptOptions<Functions> = {
repeatPenalty: {
lastTokens: 24,
penalty: 1.12,
penalizeNewLine: true,
frequencyPenalty: 0.02,
presencePenalty: 0.02,
punishTokensFilter: (tokens: Token[]) => {
return tokens.filter((token) => {
const text = model.detokenize([token])
// allow the model to repeat tokens
// that contain the word "better"
return !text.toLowerCase().includes("@")
// TODO: Exclude usernames
})
},
},
temperature: 0.6,
signal: manager,
stopOnAbortSignal: true,
}
if (debugOutput) {
defaultOptions.onResponseChunk = (chunk) => {
const isThoughtSegment = chunk.type === "segment" &&
chunk.segmentType === "thought"
if (
chunk.type === "segment" && chunk.segmentStartTime != null
) {
printSync(` [segment start: ${chunk.segmentType}] `)
}
printSync(chunk.text)
if (chunk.type === "segment" && chunk.segmentEndTime != null) {
printSync(` [segment end: ${chunk.segmentType}] `)
}
}
}
return deepMerge(defaultOptions, options)
}
export class LLamaHandler extends BaseHandler {
joinMode = false
debugLogResponses = true
systemPrompt: string
#activeModel: string
#model: LlamaModel | null = null
#context: LlamaContext | null = null
#chatWrapper: ChatWrapper | null = null
#session: LlamaChatSession | null = null
constructor(
activeModel: string,
systemPrompt: string = "You are an AI chatbot.",
) {
super("")
this.#activeModel = activeModel
this.systemPrompt = systemPrompt
}
async calculateSystemPrompt(): Promise<string> {
return this.systemPrompt
}
override async bind(bot: Bot): Promise<void> {
await super.bind(bot)
this.prefix = "@" + this.user?.username.toLowerCase()
this.#model = await llama.loadModel({
modelPath: this.#activeModel,
defaultContextFlashAttention: true,
})
this.#context = await this.#model.createContext({
flashAttention: true,
})
logger.info("Model loaded", {
batchSize: this.#context.batchSize,
contextSize: this.#context.contextSize,
})
console.log("Model loaded", {
batchSize: this.#context.batchSize,
contextSize: this.#context.contextSize,
})
this.#chatWrapper = //new Llama3ChatWrapper()
resolveChatWrapper({
bosString: this.#model
.tokens
.bosString,
filename: this.#model
.filename,
fileInfo: this.#model
.fileInfo,
tokenizer: this.#model
.tokenizer,
}) ?? new GeneralChatWrapper()
this.#session = new LlamaChatSession({
contextSequence: this.#context.getSequence(),
chatWrapper: this.#chatWrapper,
systemPrompt: await this.calculateSystemPrompt(),
})
const channels = await bot.channels
const channel = channels.find((v) => v.tag === "public") || channels[0]
if (channel) {
await bot.sendMessage(
channel.uid,
await this.#session.prompt(
"Welcome to chat, greet everyone\n",
optionsGenerator(this.#model, this.debugLogResponses),
),
)
this.#session.resetChatHistory()
}
}
override async isMatch(message: Message): Promise<boolean> {
return message.userUID !== this.user?.uid && (this.joinMode ||
trim(message?.message, " `").toLowerCase().includes(this.prefix))
}
override async handleMessage(message: Message, bot: Bot): Promise<void> {
const session = this.#session
const user = this.user
if (!session || !user) {
return
}
let response = await session.prompt(
`@${message.username}: ${message.message}`,
optionsGenerator(this.#model!, this.debugLogResponses),
)
response = response.replace(/<think>.*?<\/think>/gs, "")
let lwResponse = response.toLowerCase()
if (lwResponse.startsWith("ai")) {
response = response.substring(2).trim()
lwResponse = response.toLowerCase()
}
if (lwResponse.startsWith(`@${user.username.toLowerCase()}`)) {
response = response.substring(user.username.length + 1).trim()
lwResponse = response.toLowerCase()
}
if (lwResponse.startsWith(`@${message.username.toLowerCase()}`)) {
response = response.substring(message.username.length + 1).trim()
lwResponse = response.toLowerCase()
}
response = trimStart(response, ":").trim()
bot.send("send_message", message.channelUID, response)
}
}

12
src/msg-handlers/ping-handler.ts Normal file
View File

@ -0,0 +1,12 @@
import { BaseHandler } from "./base-handler.ts"
import { Bot, Message } from "../snek/snek-socket.ts"
export class PingHandler extends BaseHandler {
constructor() {
super("ping")
}
override handleMessage(message: Message, bot: Bot): void {
message.reply("pong")
}
}

295
src/snek/snek-socket.ts Normal file
View File

@ -0,0 +1,295 @@
import {EventEmitter} from "node:events"
import {debounce} from "npm:lodash-es"
import {getLogger} from "@logtape/logtape"
const logger = getLogger(["ws-socket"])
export const BASE_URL = "https://snek.molodetz.nl/login.json"
const baseRequest = async (
action: string,
username: string,
password: string,
) => {
const loginPayloadRed = await fetch(BASE_URL)
console.log(loginPayloadRed)
const basePayload = await loginPayloadRed.json()
basePayload.fields.username.value = username
basePayload.fields.password.value = password
return await fetch(BASE_URL, {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
action,
form: basePayload,
}),
})
}
export const login = async (username: string, password: string) => {
const loginResult = await baseRequest("submit", username, password)
const loginPayload = await loginResult.json()
if (!("redirect_url" in loginPayload)) {
return null
}
const cookies = loginResult.headers.getSetCookie()
const authCookie =
cookies.filter((v) => v.includes("AIOHTTP_SESSION")).flatMap((v) => {
try {
return v.matchAll(
/(AIOHTTP_SESSION="?.*?"?)(?:;|^)/g,
)?.next()?.value?.[1]
} catch (e) {
return undefined
}
})[0]
return authCookie || null
}
export const validate = async (username: string, password: string) => {
return await baseRequest("validate", username, password)
}
export interface User {
color: string
created_at: string
email?: string | null
last_ping: string | null
nick: string
uid: string
updated_at: string
username: string
}
export interface Channel {
name: string
uid: string
is_moderator: boolean
is_read_only: boolean
tag: string
}
export class Message {
message: string
html: string
userUID: string
color: string
channelUID: string
createdAt: string
updatedAt: string | null
username: string
uid: string
userNick: string
bot: Bot
constructor(data: any, bot: Bot) {
this.message = data.message
this.html = data.html
this.userUID = data.user_uid
this.color = data.color
this.channelUID = data.channel_uid
this.createdAt = data.created_at
this.updatedAt = data.updated_at
this.username = data.username
this.uid = data.uid
this.userNick = data.user_nick
this.bot = bot
}
reply(message: string) {
return this.bot.send("send_message", this.channelUID, message)
}
}
export class ConnectionClosedError extends Error {
constructor(code: number, reason: string) {
super(`Connection closed with code ${code}: ${reason}`)
}
}
export class Bot extends EventEmitter<
{ message: Message[]; close: { code: number; reason: string }[] }
> {
readonly #username: string
readonly #password: string
readonly #url: string | URL
#ws: WebSocket | null = null
#processingMessages = new Map<
string,
{
data: PromiseWithResolvers<unknown>
when: Date
req: {
name: string
args: unknown[]
callId: string
}
}
>()
autoReconnect = true
get channels() {
return this.send<Channel[]>("get_channels")
}
get user() {
return this.send<User>("get_user", null)
}
get messages() {
return this.send("get_messages", null)
}
#authCookie: string | null = null
get authCookie() {
return this.#authCookie
? Promise.resolve(this.#authCookie)
: login(this.#username, this.#password).then((
cookie,
) => (this.#authCookie = cookie))
}
constructor(
username: string,
password: string,
url: string | URL = "wss://snek.molodetz.nl/rpc.ws",
) {
super()
this.#username = username
this.#password = password
this.#url = url
}
send<T>(name: string, ...args: unknown[]) {
const ws = this.#ws
if (ws && ws.readyState === ws.OPEN) {
const callId = Math.random().toString(36).slice(2)
logger.debug("Sending message", { name, args, callId })
const res = Promise.withResolvers<T>()
this.#processingMessages.set(callId, {
data: res as PromiseWithResolvers<unknown>,
when: new Date(),
req: { name, args, callId },
})
ws.send(JSON.stringify({
method: name,
args: args,
callId,
}))
return res.promise.catch((e: unknown): typeof res.promise => {
logger.error("Error sending message, retrying", {
name,
args,
callId,
e,
})
return new Promise((resolve, reject) => {
setTimeout(() => {
resolve(this.connect().then(() => this.send(name, ...args)))
}, 1000)
})
})
}
return Promise.reject(new Error("Connection not open"))
}
receive(data: string) {
const parsedData = JSON.parse(data)
const callId = parsedData.callId
const message = this.#processingMessages.get(callId)
if (message) {
logger.debug("Resolving message", {
data: parsedData.data,
callId,
raw: parsedData,
})
message.data.resolve(parsedData.data)
this.#processingMessages.delete(callId)
} else {
logger.debug("Emitting message", parsedData)
this.emit("message", new Message(parsedData, this))
}
}
connect = debounce(
function (this: Bot) {
if (this.#ws) {
return this.send("get_user", null)
}
logger.debug("Connecting to", { url: this.#url })
const connectedPromise = Promise.withResolvers<unknown>()
this.#ws = new WebSocket(this.#url)
this.#ws.onopen = () => {
logger.debug("Connected")
connectedPromise.resolve(
this.send("login", this.#username, this.#password),
)
}
this.#ws.onmessage = (event) => {
logger.debug("Received message", { data: event.data })
this.receive(event.data)
}
this.#ws.onclose = (event) => {
logger.warn("Connection closed", {
code: event.code,
reason: event.reason,
})
this.#ws = null
this.emit("close", { code: event.code, reason: event.reason })
this.#processingMessages.forEach((message) => {
message.data.reject(
new ConnectionClosedError(event.code, event.reason),
)
})
this.#processingMessages.clear()
if (this.autoReconnect) {
this.connect()
}
}
this.#ws.onerror = (event) => {
logger.error("Connection error", { event })
}
return connectedPromise.promise
},
100,
{ leading: true, trailing: false },
)
sendMessage(channelUID: string, message: string) {
return this.send("send_message", channelUID, message)
}
async uploadFiles(channelUID: string, ...files: File[]) {
const imageForm = new FormData()
imageForm.append("channel_uid", channelUID)
files.forEach((file) => {
imageForm.append("files[]", file)
})
return await fetch("https://snek.molodetz.nl/drive.bin", {
method: "POST",
body: imageForm,
headers: {
"Cookie": await this.authCookie,
},
})
}
}

133
src/util/img-gen.ts Normal file
View File

@ -0,0 +1,133 @@
import {randomUUID} from "node:crypto"
import {getLogger} from "@logtape/logtape"
const logger = getLogger(["img-gen"])
export class ImgGen {
#host: string
#clientId = randomUUID()
#ws: WebSocket
#promptQueue = new Map<
string,
{ promise: PromiseWithResolvers<unknown[]>; msgs: unknown[] }
>()
constructor(host: string = "127.0.0.1:8188") {
this.#host = host
this.#initiateWebSocket()
}
async queuePrompt(prompt: string) {
const response = await fetch(`http://${this.#host}/prompt`, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ prompt, client_id: this.#clientId }),
})
return await response.json()
}
async getImage(filename: string, subfolder: string, folderType: string) {
const response = await fetch(
`http://${this.#host}/view?${
new URLSearchParams({ filename, subfolder, type: folderType })
.toString()
}`,
)
return await response.arrayBuffer()
}
async getHistory(promptId: string) {
const response = await fetch(`http://${this.#host}/history/${promptId}`)
return await response.json()
}
#initiateWebSocket() {
if (
!this.#ws || this.#ws.readyState === this.#ws.CLOSED ||
this.#ws.readyState === this.#ws.CLOSING
) {
const res = Promise.withResolvers<void>()
this.#ws = new WebSocket(
`ws://${this.#host}/ws?clientId=${this.#clientId}`,
)
let lastExportId: string | null = null
this.#ws.addEventListener("message", (event) => {
if (typeof event.data === "string") {
console.log("Received message", event, event.data)
const data = JSON.parse(event.data)
const res = this.#promptQueue.get(data.data.prompt_id)
if (res) {
res.msgs.push(data)
if (data.type === "execution_success") {
res.promise.resolve(res.msgs)
this.#promptQueue.delete(data.data.prompt_id)
} else if (data.type === "executing") {
if (data.data.node === "save_image_websocket_node") {
lastExportId = data.data.prompt_id
}
}
}
} else {
console.error(
"Received img message",
event,
"assuming is part of",
lastExportId,
)
if (lastExportId) {
const res = this.#promptQueue.get(lastExportId)
if (res) {
res.msgs.push(event.data)
}
}
}
})
this.#ws.addEventListener("open", () => {
res.resolve()
}, { once: true })
this.#ws.addEventListener("close", () => {
logger.error("WebSocket closed")
for (const [_, { promise }] of this.#promptQueue) {
promise.reject(new Error("WebSocket closed"))
}
this.#promptQueue.clear()
this.#ws = null
this.#initiateWebSocket()
})
this.#ws.addEventListener("error", (e) => {
logger.error("WebSocket error", {e})
})
return res.promise
} else if (this.#ws.readyState === this.#ws.CONNECTING) {
return new Promise<void>((resolve) => {
this.#ws.addEventListener("open", () => {
resolve()
}, { once: true })
})
}
return Promise.resolve()
}
async dispatchPrompt(prompt: Record<string, unknown> | string) {
const res = Promise.withResolvers<unknown[]>()
await this.#initiateWebSocket()
if (this.#ws.readyState !== this.#ws.OPEN) {
return Promise.reject(new Error("WebSocket not open"))
}
const promptData = await this.queuePrompt(prompt as string)
console.log("Prompt data", promptData)
this.#promptQueue.set(promptData.prompt_id, { promise: res, msgs: [] })
return res.promise
}
}

25
src/util/logging.ts Normal file
View File

@ -0,0 +1,25 @@
import {configure, getConsoleSink} from "@logtape/logtape";
await configure({
sinks: {
console: getConsoleSink(
{
formatter: (logEvent) => {
const { timestamp, level, category, message, properties } = logEvent
return `${timestamp} [${level.toUpperCase()}] [${category}] ${message} ${
properties && Object.keys(properties).length
? JSON.stringify(properties)
: ""
}`
},
},
),
},
loggers: [
{ category: "ai-app", lowestLevel: "debug", sinks: ["console"] },
{ category: "img-gen", lowestLevel: "debug", sinks: ["console"] },
{ category: "img-gen-handler", lowestLevel: "debug", sinks: ["console"] },
{ category: "llama-gen-handler", lowestLevel: "debug", sinks: ["console"] },
{ category: "ws-socket", lowestLevel: "debug", sinks: ["console"] },
],
})

19
src/ws-snek-image-bot.ts Normal file
View File

@ -0,0 +1,19 @@
import "@std/dotenv/load"
import "./util/logging.ts"
import {Bot} from "./snek/snek-socket.ts"
import {PingHandler} from "./msg-handlers/ping-handler.ts";
import {ImgGenHandler} from "./msg-handlers/img-gen-handler.ts";
const bot = new Bot(
Deno.env.get("SNEK_USERNAME")!,
Deno.env.get("SNEK_PASSWORD")!,
)
await bot.connect()
const user = await bot.user
console.log("We are user: ", user, await bot.authCookie)
new PingHandler().bind(bot)
new ImgGenHandler().bind(bot)

24
src/ws-snek-llama-bot.ts Normal file
View File

@ -0,0 +1,24 @@
import "@std/dotenv/load"
import "./util/logging.ts"
import {Bot} from "./snek/snek-socket.ts"
import {PingHandler} from "./msg-handlers/ping-handler.ts";
import {LLamaHandler} from "./msg-handlers/llama-handler.ts";
const bot = new Bot(
Deno.env.get("SNEK_USERNAME")!,
Deno.env.get("SNEK_PASSWORD")!,
)
await bot.connect()
const user = await bot.user
console.log("We are user: ", user, await bot.authCookie)
await Promise.all([
new PingHandler().bind(bot),
new LLamaHandler(
Deno.env.get("SNEK_LLAMA_MODEL")!,
Deno.env.get("SNEK_LLAMA_SYSTEM_PROMPT")!,
).bind(bot),
])

199
templates/img-gen/SD3.5-large-turbo.eta Normal file
View File

@ -0,0 +1,199 @@
{
"4": {
"inputs": {
"ckpt_name": "stableDiffusion35_largeTurbo.safetensors"
},
"class_type": "CheckpointLoaderSimple",
"_meta": {
"title": "Load Checkpoint"
}
},
"6": {
"inputs": {
"text": "<%= it.prompt %>",
"clip": [
"11",
0
]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"8": {
"inputs": {
"samples": [
"294",
0
],
"vae": [
"4",
2
]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAE Decode"
}
},
"11": {
"inputs": {
"clip_name1": "clip_g.safetensors",
"clip_name2": "clip_l.safetensors",
"clip_name3": "t5xxl_fp16.safetensors"
},
"class_type": "TripleCLIPLoader",
"_meta": {
"title": "TripleCLIPLoader"
}
},
"13": {
"inputs": {
"shift": 3,
"model": [
"4",
0
]
},
"class_type": "ModelSamplingSD3",
"_meta": {
"title": "ModelSamplingSD3"
}
},
"67": {
"inputs": {
"conditioning": [
"71",
0
]
},
"class_type": "ConditioningZeroOut",
"_meta": {
"title": "ConditioningZeroOut"
}
},
"68": {
"inputs": {
"start": 0.1,
"end": 1,
"conditioning": [
"67",
0
]
},
"class_type": "ConditioningSetTimestepRange",
"_meta": {
"title": "ConditioningSetTimestepRange"
}
},
"69": {
"inputs": {
"conditioning_1": [
"68",
0
],
"conditioning_2": [
"70",
0
]
},
"class_type": "ConditioningCombine",
"_meta": {
"title": "Conditioning (Combine)"
}
},
"70": {
"inputs": {
"start": 0,
"end": 0.1,
"conditioning": [
"71",
0
]
},
"class_type": "ConditioningSetTimestepRange",
"_meta": {
"title": "ConditioningSetTimestepRange"
}
},
"71": {
"inputs": {
"text": "<%= it.negativePrompt ||'' %>",
"clip": [
"11",
0
]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"135": {
"inputs": {
"width": <%= it.width || 1024 %>,
"height": <%= it.height || 1024 %>,
"batch_size": <%= it.batchSize || 1 %>
},
"class_type": "EmptySD3LatentImage",
"_meta": {
"title": "EmptySD3LatentImage"
}
},
"294": {
"inputs": {
"seed": <%= it.seed || it.randomSeed() %>,
"steps": <%= it.steps || 2 %>,
"cfg": <%= it.cfg || 1 %>,
"sampler_name": "<%= it.sampler || "euler" %>",
"scheduler": "beta",
"denoise": 1,
"model": [
"13",
0
],
"positive": [
"6",
0
],
"negative": [
"69",
0
],
"latent_image": [
"135",
0
]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"301": {
"inputs": {
"filename_prefix": "ComfyUI",
"images": [
"8",
0
]
},
"class_type": "SaveImage",
"_meta": {
"title": "Save Image"
}
},
"save_image_websocket_node": {
"inputs": {
"images": [
"8",
0
]
},
"class_type": "SaveImageWebsocket",
"_meta": {
"title": "SaveImageWebsocket"
}
}
}

199
templates/img-gen/SD3.5-large.eta Normal file
View File

@ -0,0 +1,199 @@
{
"4": {
"inputs": {
"ckpt_name": "stableDiffusion35_large.safetensors"
},
"class_type": "CheckpointLoaderSimple",
"_meta": {
"title": "Load Checkpoint"
}
},
"6": {
"inputs": {
"text": "<%= it.prompt %>",
"clip": [
"11",
0
]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"8": {
"inputs": {
"samples": [
"294",
0
],
"vae": [
"4",
2
]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAE Decode"
}
},
"11": {
"inputs": {
"clip_name1": "clip_g.safetensors",
"clip_name2": "clip_l.safetensors",
"clip_name3": "t5xxl_fp16.safetensors"
},
"class_type": "TripleCLIPLoader",
"_meta": {
"title": "TripleCLIPLoader"
}
},
"13": {
"inputs": {
"shift": 3,
"model": [
"4",
0
]
},
"class_type": "ModelSamplingSD3",
"_meta": {
"title": "ModelSamplingSD3"
}
},
"67": {
"inputs": {
"conditioning": [
"71",
0
]
},
"class_type": "ConditioningZeroOut",
"_meta": {
"title": "ConditioningZeroOut"
}
},
"68": {
"inputs": {
"start": 0.1,
"end": 1,
"conditioning": [
"67",
0
]
},
"class_type": "ConditioningSetTimestepRange",
"_meta": {
"title": "ConditioningSetTimestepRange"
}
},
"69": {
"inputs": {
"conditioning_1": [
"68",
0
],
"conditioning_2": [
"70",
0
]
},
"class_type": "ConditioningCombine",
"_meta": {
"title": "Conditioning (Combine)"
}
},
"70": {
"inputs": {
"start": 0,
"end": 0.1,
"conditioning": [
"71",
0
]
},
"class_type": "ConditioningSetTimestepRange",
"_meta": {
"title": "ConditioningSetTimestepRange"
}
},
"71": {
"inputs": {
"text": "<%= it.negativePrompt ||'' %>",
"clip": [
"11",
0
]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"135": {
"inputs": {
"width": <%= it.width || 1024 %>,
"height": <%= it.height || 1024 %>,
"batch_size": <%= it.batchSize || 1 %>
},
"class_type": "EmptySD3LatentImage",
"_meta": {
"title": "EmptySD3LatentImage"
}
},
"294": {
"inputs": {
"seed": <%= it.seed || it.randomSeed() %>,
"steps": <%= it.steps || 2 %>,
"cfg": <%= it.cfg || 1 %>,
"sampler_name": "<%= it.sampler || "euler" %>",
"scheduler": "beta",
"denoise": 1,
"model": [
"13",
0
],
"positive": [
"6",
0
],
"negative": [
"69",
0
],
"latent_image": [
"135",
0
]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"301": {
"inputs": {
"filename_prefix": "ComfyUI",
"images": [
"8",
0
]
},
"class_type": "SaveImage",
"_meta": {
"title": "Save Image"
}
},
"save_image_websocket_node": {
"inputs": {
"images": [
"8",
0
]
},
"class_type": "SaveImageWebsocket",
"_meta": {
"title": "SaveImageWebsocket"
}
}
}

204
templates/img-gen/flux-dev.eta Normal file
View File

@ -0,0 +1,204 @@
{
"5": {
"inputs": {
"width": <%= it.width || 1024 %>,
"height": <%= it.height || 1024 %>,
"batch_size": <%= it.batchSize || 1 %>
},
"class_type": "EmptyLatentImage",
"_meta": {
"title": "Empty Latent Image"
}
},
"6": {
"inputs": {
"text": "<%= it.prompt ||'' %>",
"clip": [
"11",
0
]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"8": {
"inputs": {
"samples": [
"13",
0
],
"vae": [
"10",
0
]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAE Decode"
}
},
"9": {
"inputs": {
"filename_prefix": "MarkuryFLUX",
"images": [
"8",
0
]
},
"class_type": "SaveImage",
"_meta": {
"title": "Save Image"
}
},
"10": {
"inputs": {
"vae_name": "ae.safetensors"
},
"class_type": "VAELoader",
"_meta": {
"title": "Load VAE"
}
},
"11": {
"inputs": {
"clip_name1": "t5xxl_fp16.safetensors",
"clip_name2": "clip_l.safetensors",
"type": "flux",
"device": "default"
},
"class_type": "DualCLIPLoader",
"_meta": {
"title": "DualCLIPLoader"
}
},
"12": {
"inputs": {
"unet_name": "flux_dev.safetensors",
"weight_dtype": "fp8_e4m3fn"
},
"class_type": "UNETLoader",
"_meta": {
"title": "Load Diffusion Model"
}
},
"13": {
"inputs": {
"noise": [
"25",
0
],
"guider": [
"22",
0
],
"sampler": [
"16",
0
],
"sigmas": [
"17",
0
],
"latent_image": [
"5",
0
]
},
"class_type": "SamplerCustomAdvanced",
"_meta": {
"title": "SamplerCustomAdvanced"
}
},
"16": {
"inputs": {
"sampler_name": "<%= it.sampler || "euler" %>"
},
"class_type": "KSamplerSelect",
"_meta": {
"title": "KSamplerSelect"
}
},
"17": {
"inputs": {
"scheduler": "beta",
"steps": <%= it.steps || 30 %>,
"denoise": 1,
"model": [
"61",
0
]
},
"class_type": "BasicScheduler",
"_meta": {
"title": "BasicScheduler"
}
},
"22": {
"inputs": {
"model": [
"61",
0
],
"conditioning": [
"60",
0
]
},
"class_type": "BasicGuider",
"_meta": {
"title": "BasicGuider"
}
},
"25": {
"inputs": {
"noise_seed": <%= it.seed || it.randomSeed() %>
},
"class_type": "RandomNoise",
"_meta": {
"title": "RandomNoise"
}
},
"60": {
"inputs": {
"guidance": 4,
"conditioning": [
"6",
0
]
},
"class_type": "FluxGuidance",
"_meta": {
"title": "FluxGuidance"
}
},
"61": {
"inputs": {
"max_shift": 1.15,
"base_shift": 0.5,
"width": <%= it.width || 1024 %>,
"height": <%= it.height || 1024 %>,
"model": [
"12",
0
]
},
"class_type": "ModelSamplingFlux",
"_meta": {
"title": "ModelSamplingFlux"
}
},
"save_image_websocket_node": {
"inputs": {
"images": [
"8",
0
]
},
"class_type": "SaveImageWebsocket",
"_meta": {
"title": "SaveImageWebsocket"
}
}
}

149
templates/img-gen/mochi.eta Normal file
View File

@ -0,0 +1,149 @@
{
"3": {
"inputs": {
"seed": <%= it.seed || it.randomSeed() %>,
"steps": <%= it.steps || 30 %>,
"cfg": <%= it.cfg || 4.5 %>,
"sampler_name": "<%= it.sampler || "euler" %>",
"scheduler": "beta",
"denoise": 1,
"model": [
"37",
0
],
"positive": [
"6",
0
],
"negative": [
"7",
0
],
"latent_image": [
"21",
0
]
},
"class_type": "KSampler",
"_meta": {
"title": "KSampler"
}
},
"6": {
"inputs": {
"text": "<%= it.prompt %>",
"clip": [
"38",
0
]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"7": {
"inputs": {
"text": "<%= it.negativePrompt || '' %>",
"clip": [
"38",
0
]
},
"class_type": "CLIPTextEncode",
"_meta": {
"title": "CLIP Text Encode (Prompt)"
}
},
"8": {
"inputs": {
"samples": [
"3",
0
],
"vae": [
"39",
0
]
},
"class_type": "VAEDecode",
"_meta": {
"title": "VAE Decode"
}
},
"21": {
"inputs": {
"width": <%= it.width || 848 %>,
"height": <%= it.height || 480 %>,
"length": <%= it.frames || 37 %>,
"batch_size": <%= it.batchSize || 1 %>
},
"class_type": "EmptyMochiLatentVideo",
"_meta": {
"title": "EmptyMochiLatentVideo"
}
},
"28": {
"inputs": {
"filename_prefix": "ComfyUI",
"fps": 24,
"lossless": false,
"quality": 80,
"method": "default",
"images": [
"8",
0
]
},
"class_type": "SaveAnimatedWEBP",
"_meta": {
"title": "SaveAnimatedWEBP"
}
},
"37": {
"inputs": {
"unet_name": "mochi_preview_bf16.safetensors",
"weight_dtype": "default"
},
"class_type": "UNETLoader",
"_meta": {
"title": "Load Diffusion Model"
}
},
"38": {
"inputs": {
"clip_name": "t5xxl_fp16.safetensors",
"type": "mochi",
"device": "default"
},
"class_type": "CLIPLoader",
"_meta": {
"title": "Load CLIP"
}
},
"39": {
"inputs": {
"vae_name": "mochi_vae.safetensors"
},
"class_type": "VAELoader",
"_meta": {
"title": "Load VAE"
}
},
"save_image_websocket_node": {
"inputs": {
"fps": 24,
"lossless": false,
"quality": 80,
"method": "default",
"images": [
"8",
0
]
},
"class_type": "SaveAnimatedWEBPWebsocket",
"_meta": {
"title": "SaveAnimatedWEBPWebsocket"
}
}
}