comfyui-image-generator

import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { readFile, writeFile } from "node:fs/promises";
import { join, basename } from "node:path";

const COMFYUI_URL = "http://127.0.0.1:8188";

const WORKFLOW_PATH =
  String.raw`C:\Users\Kirill\Downloads\StabilityMatrix\lmstudiotool\DefaultImage.json`;

const PROMPT_NODE_ID = "3";

type ComfyImageResult = {
  filename: string;
  subfolder: string;
  type: string;
};

async function loadWorkflow(): Promise<any> {
  const raw = await readFile(WORKFLOW_PATH, "utf-8");
  return JSON.parse(raw);
}

async function sendPrompt(promptText: string): Promise<string> {
  const workflow = await loadWorkflow();

  workflow[PROMPT_NODE_ID].inputs.text = promptText;

  const response = await fetch(`${COMFYUI_URL}/prompt`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json"
    },
    body: JSON.stringify({
      prompt: workflow,
      client_id: crypto.randomUUID()
    })
  });

  if (!response.ok) {
    throw new Error(`ComfyUI /prompt failed: ${response.status} ${response.statusText}`);
  }

  const result: any = await response.json();
  return result.prompt_id;
}

async function getHistory(promptId: string): Promise<any> {
  const response = await fetch(`${COMFYUI_URL}/history/${promptId}`);

  if (!response.ok) {
    return null;
  }

  return response.json();
}

async function waitForImage(promptId: string, timeoutMs = 180_000): Promise<ComfyImageResult> {
  const startedAt = Date.now();

  while (true) {
    if (Date.now() - startedAt > timeoutMs) {
      throw new Error("Timeout waiting for ComfyUI image generation");
    }

    const history = await getHistory(promptId);

    if (history?.[promptId]) {
      const outputs = history[promptId].outputs ?? {};

      for (const nodeOutput of Object.values(outputs) as any[]) {
        const images = nodeOutput.images;

        if (images && images.length > 0) {
          const image = images[0];

          return {
            filename: image.filename,
            subfolder: image.subfolder ?? "",
            type: image.type ?? "output"
          };
        }
      }
    }

    await new Promise((resolve) => setTimeout(resolve, 1000));
  }
}

function buildComfyUIImageUrl(image: ComfyImageResult): string {
  const params = new URLSearchParams({
    filename: image.filename,
    subfolder: image.subfolder,
    type: image.type
  });

  return `${COMFYUI_URL}/view?${params.toString()}`;
}

async function downloadImageToLmStudioWorkingDir(
  imageUrl: string,
  originalFilename: string,
  workingDirectory: string
): Promise<string> {
  const response = await fetch(imageUrl);

  if (!response.ok) {
    throw new Error(`Failed to download image from ComfyUI: ${response.status} ${response.statusText}`);
  }

  const bytes = new Uint8Array(await response.arrayBuffer());

  const safeFilename = `${Date.now()}-${basename(originalFilename)}`;
  const filePath = join(workingDirectory, safeFilename);

  await writeFile(filePath, bytes);

  // Важно: именно такой формат лучше воспринимается LM Studio.
  // Пример из рабочих plugins: заменить \ на / и убрать C:
  const lmStudioPath = filePath
    .replaceAll('\\', "/")
    .replace(/^C:/, "");

  return lmStudioPath;
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const generateComfyUIImage = tool({
    name: "generate_comfyui_image",
    description:
      "Generate an image with ComfyUI from a prompt and display it inside LM Studio chat. Use this when the user asks to create or generate an image.",
    parameters: {
      prompt: z.string().describe("The image prompt. Prefer English prompts for best image generation results.")
    },
    implementation: async ({ prompt }, { status }) => {
      status("Sending prompt to ComfyUI...");

      const promptId = await sendPrompt(prompt);

      status("Waiting for ComfyUI image generation...");

      const image = await waitForImage(promptId);

      status("Downloading generated image into LM Studio working directory...");

      const imageUrl = buildComfyUIImageUrl(image);
      const workingDirectory = ctl.getWorkingDirectory();

      const lmStudioImagePath = await downloadImageToLmStudioWorkingDir(
        imageUrl,
        image.filename,
        workingDirectory
      );

      status("Image generated successfully.");

      return `![Generated image](${lmStudioImagePath})`;
    }
  });

  return [generateComfyUIImage];
}