src / toolsProvider.ts
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 ``;
}
});
return [generateComfyUIImage];
}