Merge pull request #161 from harmony-one/sd_images_mul_models_2

DRAFT: SD images multiple models V2

src/modules/sd-images/SDImagesBotBase.ts

@@ -0,0 +1,65 @@
+import { SDNodeApi, getModelByParam, IModel } from "./api";
+import { OnMessageContext, OnCallBackQueryData } from "../types";
+import { sleep, uuidv4 } from "./utils";
+import { InlineKeyboard, InputFile } from "grammy";
+
+export class SDImagesBotBase {
+    sdNodeApi: SDNodeApi;
+
+    queue: string[] = [];
+
+    constructor() {
+        this.sdNodeApi = new SDNodeApi();
+    }
+
+    waitingQueue = async (uuid: string, ctx: OnMessageContext | OnCallBackQueryData,) => {
+        this.queue.push(uuid);
+
+        let idx = this.queue.findIndex((v) => v === uuid);
+
+        if (idx !== 0) {
+            ctx.reply(
+                `You are ${idx + 1}/${this.queue.length
+                }, wait about ${idx * 30} seconds`
+            );
+        }
+
+        // waiting queue
+        while (idx !== 0) {
+            await sleep(3000 * this.queue.findIndex((v) => v === uuid));
+            idx = this.queue.findIndex((v) => v === uuid);
+        }
+    }
+
+    generateImage = async (
+        ctx: OnMessageContext | OnCallBackQueryData,
+        refundCallback: (reason?: string) => void,
+        prompt: string,
+        model: IModel,
+        seed?: number
+    ) => {
+        const uuid = uuidv4();
+
+        try {
+            await this.waitingQueue(uuid, ctx);
+
+            ctx.chatAction = "upload_photo";
+
+            const imageBuffer = await this.sdNodeApi.generateImage(
+                prompt,
+                model,
+                seed
+            );
+
+            await ctx.replyWithPhoto(new InputFile(imageBuffer), {
+                caption: `/${model.aliases[0]} ${prompt}`,
+            });
+        } catch (e) {
+            console.error(e);
+            ctx.reply(`Error: something went wrong... Refunding payments`);
+            refundCallback();
+        }
+
+        this.queue = this.queue.filter((v) => v !== uuid);
+    }
+}

src/modules/sd-images/api/helpers.ts

@@ -0,0 +1,8 @@
+export const waitingExecute = (fn: () => Promise<any>, ms: number) => new Promise((resolve, reject) => {
+    const timeoutId = setTimeout(() => { 
+      console.error('SD images Error: waitingExecute time is up');
+      reject('Error: waitingExecute time is up');
+    }, ms);
+
+    fn().then(resolve).catch(reject).finally(() => clearTimeout(timeoutId))
+  });

src/modules/sd-images/api/index.ts

@@ -0,0 +1,2 @@
+export * from './sd-node-api';
+export * from './models-config';

src/modules/sd-images/api/models-config.ts

@@ -0,0 +1,135 @@
+export interface IModel {
+    path: string;
+    name: string;
+    id: string;
+    hash: string;
+    shortName: string;
+    link: string;
+    baseModel: 'SD 1.5' | 'SDXL 1.0';
+    aliases: string[];
+    defaultPrompt: string;
+}
+
+export const MODELS_CONFIGS: IModel[] = [
+    {
+        path: "deliberate_v2.safetensors",
+        name: "Deliberate",
+        id: '4823',
+        hash: '9ABA26ABDF',
+        shortName: 'deliberate',
+        link: 'https://civitai.com/models/4823/deliberate',
+        baseModel: 'SD 1.5',
+        aliases: ['del', '4823', 'h9ab'],
+        defaultPrompt: 'a cute kitten made out of metal, (cyborg:1.1), ([tail | detailed wire]:1.3), (intricate details), hdr, (intricate details, hyperdetailed:1.2), cinematic shot, vignette, centered',
+    },
+    {
+        path: "dreamshaper_8.safetensors",
+        name: "DreamShaper",
+        id: '4384',
+        hash: '879DB523C3',
+        shortName: 'dreamshaper',
+        link: 'https://civitai.com/models/4384/dreamshaper',
+        baseModel: 'SD 1.5',
+        aliases: ['dream', '4384', 'h879'],
+        defaultPrompt: '8k portrait of beautiful cyborg with brown hair, intricate, elegant, highly detailed, majestic, digital photography, art by artgerm and ruan jia and greg rutkowski surreal painting gold butterfly filigree, broken glass, (masterpiece, sidelighting, finely detailed beautiful eyes: 1.2), hdr, <lora:more_details:0.36>'
+    },
+    {
+        path: "majicmixRealistic_betterV2V25.safetensors",
+        name: 'majicMIX realistic 麦橘写实',
+        id: '43331',
+        hash: 'D7E2AC2F4A',
+        shortName: 'majicmix-realistic',
+        link: 'https://civitai.com/models/43331/majicmix-realistic',
+        baseModel: 'SD 1.5',
+        aliases: ['maji', '4333', 'h7e2'],
+        defaultPrompt: '1girl,sitting on a cozy couch,crossing legs,soft light'
+    },
+    {
+        path: "revAnimated_v122.safetensors",
+        name: 'ReV Animated',
+        id: '7371',
+        hash: '4199BCDD14',
+        shortName: 'rev-animated',
+        link: 'https://civitai.com/models/7371/rev-animated',
+        baseModel: 'SD 1.5',
+        aliases: ['rev', '7371', 'h419'],
+        defaultPrompt: '((best quality)), ((masterpiece)), (detailed), alluring succubus, ethereal beauty, perched on a cloud, (fantasy illustration:1.3), enchanting gaze, captivating pose, delicate wings, otherworldly charm, mystical sky, (Luis Royo:1.2), (Yoshitaka Amano:1.1), moonlit night, soft colors, (detailed cloudscape:1.3), (high-resolution:1.2)'
+    },
+    // {
+    //     path: "v1-5-pruned-emaonly.safetensors",
+    //     name: '',
+    //     id: '',
+    //     hash: '',
+    //     shortName: '',
+    //     link: '',
+    //     baseModel: 'SD 1.5'
+    // },
+    {
+        path: "animePastelDream_softBakedVae.safetensors",
+        name: 'Anime Pastel Dream',
+        id: '23521',
+        hash: '4BE38C1A17',
+        shortName: 'anime-pastel-dream',
+        link: 'https://civitai.com/models/23521/anime-pastel-dream',
+        baseModel: 'SD 1.5',
+        aliases: ['anime', '2352', 'h4be'],
+        defaultPrompt: 'masterpiece, best quality, ultra-detailed, illustration,(1girl),beautiful detailed eyes, looking at viewer, close up, (breast focus), pink hair, shy, cat ears'
+    },
+    {
+        path: 'cyberrealistic_v33.safetensors',
+        name: 'CyberRealistic',
+        id: '15003',
+        hash: '7A4DBBA12F',
+        shortName: 'cyberrealistic',
+        link: 'https://civitai.com/models/15003/cyberrealistic',
+        baseModel: 'SD 1.5',
+        aliases: ['cyber', '1500', 'h7a4'],
+        defaultPrompt: '<lora:LowRA:0.6> (8k, RAW photo, highest quality), beautiful girl, close up, dress, (detailed eyes:0.8), defiance512, (looking at the camera:1.4), (highest quality), (best shadow), intricate details, interior, ginger hair:1.3, dark studio, muted colors, freckles   <lora:epiNoiseoffset_v2Pynoise:1.2>',
+    },
+    {
+        path: 'dreamshaperXL10_alpha2Xl10.safetensors',
+        name: 'DreamShaper XL1.0',
+        id: '112902',
+        hash: '0F1B80CFE8',
+        shortName: 'dreamshaper-xl10',
+        link: 'https://civitai.com/models/112902/dreamshaper-xl10',
+        baseModel: 'SDXL 1.0',
+        aliases: ['xl_dream', '1129', 'h0f1'],
+        defaultPrompt: 'barry allen the flash on wheelchair moving at supersonic speed creating flame trails, speed trails, motion blur, electricity speed outdoor, realistic highly detailed cinematic cinematography, movie shots footage,',
+    },
+    {
+        path: 'sdXL_v10VAEFix.safetensors',
+        name: 'SD XL',
+        id: '101055',
+        hash: 'E6BB9EA85B',
+        shortName: 'sd-xl',
+        link: 'https://civitai.com/models/101055/sd-xl',
+        baseModel: 'SDXL 1.0',
+        aliases: ['xl', '1010', 'he6b'],
+        defaultPrompt: 'A cat is sitting in a kimono, in the style of renaissance - inspired chiaroscuro, hyper - realistic portraiture, nicolas mignard, old master influenced fantasy, portraitures with hidden meanings, dom qwek, art of burma'
+    },
+    {
+        path: 'sdxlUnstableDiffusers_v5UnchainedSlayer.safetensors',
+        name: 'SDXL Unstable Diffusers ☛ YamerMIX',
+        id: '84040',
+        hash: 'EF924AAE79',
+        shortName: 'sdxl-unstable-diffusers-yamermix',
+        link: 'https://civitai.com/models/84040/sdxl-unstable-diffusers-yamermix',
+        baseModel: 'SDXL 1.0',
+        aliases: ['xl_dif', '8404', 'hef9'],
+        defaultPrompt: 'pastel color, from above, upper body, depth of field, masterpiece, best quality, best quality, 1girl sitting on a swing, school uniform, black hair, blue eyes, autumn, park'
+    }
+];
+
+export const getModelByParam = (param: string) => {
+    const model = MODELS_CONFIGS.find(m =>
+        m.id === param ||
+        m.hash === param ||
+        m.shortName === param ||
+        m.aliases.includes(param)
+    );
+
+    return model;
+}
+
+export const modelsAliases = [];

src/modules/sd-images/api/sd-node-api.ts

@@ -0,0 +1,55 @@
+import { Client } from './sd-node-client'
+import { MODELS_CONFIGS, getModelByParam, IModel } from './models-config';
+
+const NEGATIVE_PROMPT = '(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation';
+
+export class SDNodeApi {
+  client: Client;
+
+  constructor() {
+    this.client = new Client()
+  }
+
+  generateImage = async (prompt: string, model: IModel, seed?: number) => {
+    const { images } = await this.client.txt2img({
+      prompt,
+      negativePrompt: NEGATIVE_PROMPT,
+      width: model.baseModel === 'SDXL 1.0' ? 1024 : 512,
+      height: model.baseModel === 'SDXL 1.0' ? 1024 : 768,
+      steps: 26,
+      batchSize: 1,
+      cfgScale: 7,
+      seed,
+      model: model.path
+    })
+
+    return images[0];
+  }
+
+  generateImagesPreviews = async (prompt: string, model: IModel) => {
+    const params = {
+      prompt,
+      negativePrompt: NEGATIVE_PROMPT,
+      width: model.baseModel === 'SDXL 1.0' ? 1024 : 512,
+      height: model.baseModel === 'SDXL 1.0' ? 1024 : 768,
+      steps: 15,
+      batchSize: 1,
+      cfgScale: 7,
+      model: model.path
+    };
+
+    const res = await Promise.all([
+      this.client.txt2img(params),
+      this.client.txt2img(params),
+      this.client.txt2img(params),
+      this.client.txt2img(params)
+    ]);
+
+    return {
+      images: res.map(r => r.images[0]),
+      parameters: {},
+      all_seeds: res.map(r => r.all_seeds[0]),
+      info: ''
+    };
+  }
+}

src/modules/sd-images/sd-node-client.ts → src/modules/sd-images/api/sd-node-client.ts

@@ -1,7 +1,9 @@
-import { ComfyClient } from '../qrcode/comfy/ComfyClient';
-import config from "../../config";
-import { sleep, waitingExecute } from './utils';
+import { ComfyClient } from '../../qrcode/comfy/ComfyClient';
+import config from "../../../config";
+import { sleep } from '../utils';
 import { buildImgPrompt } from './text_to_img_config';
+import { MODELS_CONFIGS } from './models-config';
+import { waitingExecute } from './helpers';
 
 export type Txt2ImgOptions = {
     hires?: {
@@ -31,7 +33,7 @@ export type Txt2ImgOptions = {
         name: string
         args?: string[]
     }
-    model?: MODELS
+    model?: string;
 }
 
 export type Txt2ImgResponse = {
@@ -41,19 +43,16 @@ export type Txt2ImgResponse = {
     info: string
 }
 
-export enum MODELS {
-    "XL_BASE_1.0" = "sd_xl_base_1.0.safetensors",
-}
-
 const getRandomSeed = () => Math.round(Math.random() * 1e15);
 
 export class Client {
     constructor() { }
 
-    txt2img = async (options: Txt2ImgOptions): Promise<Txt2ImgResponse> => {
+    txt2img = async (options: Txt2ImgOptions, serverConfig?: { host: string, wsHost: string }): Promise<Txt2ImgResponse> => {
         const comfyClient = new ComfyClient({
             host: config.comfyHost,
-            wsHost: config.comfyWsHost
+            wsHost: config.comfyWsHost,
+            ...serverConfig
         });
 
         try {
@@ -70,12 +69,11 @@ export class Client {
                 ...options,
                 seed,
                 clientId: comfyClient.clientId,
-                model: MODELS['XL_BASE_1.0']
             });
 
             const r = await comfyClient.queuePrompt(prompt);
 
-            const promptResult = await waitingExecute(() => comfyClient.waitingPromptExecution(r.prompt_id), 1000 * 120);
+            const promptResult = await waitingExecute(() => comfyClient.waitingPromptExecution(r.prompt_id), 1000 * 180);
 
             const history = await comfyClient.history(r.prompt_id);
 
@@ -88,7 +86,7 @@ export class Client {
             return {
                 images,
                 parameters: {},
-                all_seeds: history.outputs['9'].images.map((i, idx) => String(seed + idx)),
+                all_seeds: [String(seed)],
                 info: ''
             } as Txt2ImgResponse;
         } catch (e) {

src/modules/sd-images/text_to_img_config.ts → src/modules/sd-images/api/text_to_img_config.ts

@@ -9,7 +9,7 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) {
                     seed: options.seed,
                     steps: options.steps,
                     cfg: 8,
-                    sampler_name: "euler_ancestral",
+                    sampler_name: "dpmpp_2m",
                     scheduler: "karras",
                     denoise: 1,
                     model: ["4", 0],
@@ -25,8 +25,8 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) {
             },
             5: {
                 inputs: {
-                    width: 1024,
-                    height: 1024,
+                    width: options.width || 1024,
+                    height: options.height || 1024,
                     batch_size: options.batchSize || 1,
                 },
                 class_type: "EmptyLatentImage",
@@ -35,7 +35,7 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) {
                 inputs: { text: options.prompt, clip: ["4", 1] },
                 class_type: "CLIPTextEncode",
             },
-            7: { inputs: { text: "", clip: ["4", 1] }, class_type: "CLIPTextEncode" },
+            7: { inputs: { text: options.negativePrompt, clip: ["4", 1] }, class_type: "CLIPTextEncode" },
             8: {
                 inputs: { samples: ["3", 0], vae: ["4", 2] },
                 class_type: "VAEDecode",
@@ -136,7 +136,7 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) {
                                 { name: "LATENT", type: "LATENT", links: [2], slot_index: 0 },
                             ],
                             properties: { "Node name for S&R": "EmptyLatentImage" },
-                            widgets_values: [1024, 1024, 1],
+                            widgets_values: [options.width || 1024, options.height || 1024, 1],
                         },
                         {
                             id: 9,
@@ -169,11 +169,11 @@ export function buildImgPrompt(options: Txt2ImgOptions & { clientId: string }) {
                             ],
                             properties: { "Node name for S&R": "KSampler" },
                             widgets_values: [
-                                877217196051697,
+                                options.seed,
                                 "randomize",
-                                30,
+                                options.steps,
                                 8,
-                                "euler_ancestral",
+                                "dpmpp_2m",
                                 "karras",
                                 1,
                             ],