import { ComfyUIClient } from '@ayan4m1/comfy-ui-client';
import type { Prompt } from '@ayan4m1/comfy-ui-client';

// Your prompt / workflow
const prompt: Prompt = {
  '3': {
    class_type: 'KSampler',
    inputs: {
      cfg: 8,
      denoise: 1,
      latent_image: ['5', 0],
      model: ['4', 0],
      negative: ['7', 0],
      positive: ['6', 0],
      sampler_name: 'euler',
      scheduler: 'normal',
      seed: 8566257,
      steps: 20,
    },
  },
  '4': {
    class_type: 'CheckpointLoaderSimple',
    inputs: {
      ckpt_name: 'v1-5-pruned-emaonly.cpkt',
    },
  },
  '5': {
    class_type: 'EmptyLatentImage',
    inputs: {
      batch_size: 1,
      height: 512,
      width: 512,
    },
  },
  '6': {
    class_type: 'CLIPTextEncode',
    inputs: {
      clip: ['4', 1],
      text: 'masterpiece best quality girl',
    },
  },
  '7': {
    class_type: 'CLIPTextEncode',
    inputs: {
      clip: ['4', 1],
      text: 'bad hands',
    },
  },
  '8': {
    class_type: 'VAEDecode',
    inputs: {
      samples: ['3', 0],
      vae: ['4', 2],
    },
  },
  '9': {
    class_type: 'SaveImage',
    inputs: {
      filename_prefix: 'ComfyUI',
      images: ['8', 0],
    },
  },
};

// Set the text prompt for our positive CLIPTextEncode
prompt['6'].inputs.text = 'masterpiece best quality man';

// Set the seed for our KSampler node
prompt['3'].inputs.seed = 5;

// Create client
const host = 'http://127.0.0.1:8188';
const clientId = 'baadbabe-b00b-4206-9420-deadd00d1337';
// Optional if using ComfyUI-Login
const authToken = 'b2$diovjwaer98vbu89agusd89vsu8v89s';
const client = new ComfyUIClient(host, clientId, authToken);

// Connect to server
await client.connect();

// Generate images
const images = await client.getImages(prompt);

// Save images to file
const outputDir = './tmp/output';
await client.saveImages(images, outputDir);

// Disconnect
await client.disconnect();