src / toolsProvider.ts
import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./config";
import { spawn, execSync } from "child_process";
import { promises as fs } from "fs";
import { join } from "path";
import { homedir } from "os";
function expandTilde(filepath: string): string {
if (filepath.startsWith("~/") || filepath === "~") {
return join(homedir(), filepath.slice(1));
}
return filepath;
}
export async function toolsProvider(ctl: ToolsProviderController) {
const config = ctl.getPluginConfig(configSchematics);
const tools: Tool[] = [];
const generateImageTool = tool({
name: "generate_image",
description: "Generate an image using Stable Diffusion from a text prompt. All parameters except 'prompt' are optional and have pre-configured defaults. DO NOT set optional parameters unless the user explicitly specifies them in their message. Omitting a parameter uses the user's configured default.",
parameters: {
prompt: z.string().describe("Text description of the image to generate"),
negative_prompt: z.string().optional().describe("What to avoid in the image. Omit to use default."),
seed: z.number().optional().describe("Random seed. Omit for random."),
width: z.number().optional().describe("Image width in pixels (multiple of 64, 256-1536). OMIT unless user specifies."),
height: z.number().optional().describe("Image height in pixels (multiple of 64, 256-1536). OMIT unless user specifies."),
steps: z.number().optional().describe("Sampling steps (1-50). OMIT unless user specifies."),
cfg_scale: z.number().optional().describe("CFG scale (1-20). OMIT unless user specifies."),
sampler: z.enum(["lcm", "euler_a", "euler", "dpm++2m", "dpm++2s_a", "heun"]).optional().describe("Sampling method. OMIT unless user specifies."),
scheduler: z.enum(["discrete", "karras", "exponential", "ays", "gits", "smoothstep", "sgm_uniform", "simple", "kl_optimal", "lcm", "bong_tangent"]).optional().describe("Scheduler. OMIT unless user specifies."),
upscale: z.union([z.boolean(), z.string()]).optional().describe("Whether to upscale. OMIT unless user specifies."),
upscale_factor: z.number().optional().describe("Upscale factor: 2, 3, or 4. OMIT unless user specifies.")
},
implementation: async ({ prompt, negative_prompt, seed, width: paramWidth, height: paramHeight, steps: paramSteps, cfg_scale: paramCfgScale, sampler: paramSampler, scheduler: paramScheduler, upscale: paramUpscale, upscale_factor: paramUpscaleFactorNum }) => {
// Read config values fresh on each invocation so UI changes take effect
// LLM-provided parameters override config defaults
const sdCliPath = config.get("sdCliPath");
const modelPath = config.get("modelPath");
const outputDirectory = config.get("outputDirectory");
const width = paramWidth ?? config.get("width");
const height = paramHeight ?? config.get("height");
const steps = paramSteps ?? config.get("steps");
const cfgScale = paramCfgScale ?? config.get("cfgScale");
const sampler = paramSampler ?? config.get("sampler");
const scheduler = paramScheduler ?? config.get("scheduler");
const clipOnCpu = config.get("clipOnCpu");
const vaeTiling = config.get("vaeTiling");
const enableUpscale = paramUpscale !== undefined ? (String(paramUpscale).toLowerCase() === "true" ? "true" : "false") : config.get("enableUpscale");
const upscaleModelPath = config.get("upscaleModelPath");
const upscaleTileSize = config.get("upscaleTileSize");
const upscaleFactor = paramUpscaleFactorNum !== undefined ? String(paramUpscaleFactorNum) : config.get("upscaleFactor");
// Validate configuration
if (!sdCliPath) {
return "Error: sd-cli path not configured. Please set it in the plugin settings (click on create-image in Integrations panel).";
}
if (!modelPath) {
return "Error: Model path not configured. Please set it in the plugin settings.";
}
// Resolve paths (expand ~ to home directory)
const resolvedSdCliPath = expandTilde(sdCliPath);
const resolvedModelPath = expandTilde(modelPath);
// Generate unique filename for sd-cli output
const timestamp = Date.now();
const resolvedOutputDir = expandTilde(outputDirectory || "/tmp");
const tempOutputPath = join(resolvedOutputDir, `sd_${timestamp}.png`);
// Ensure output directory exists
await fs.mkdir(resolvedOutputDir, { recursive: true });
// Default negative prompt
const negPrompt = negative_prompt || "blurry, bad quality, distorted, ugly, low resolution, watermark, text";
// Build command arguments
const args = [
"--prompt", prompt,
"--negative-prompt", negPrompt,
"--output", tempOutputPath,
"--width", String(width),
"--height", String(height),
"--steps", String(steps),
"--sampling-method", sampler,
"--scheduler", scheduler,
"--cfg-scale", String(cfgScale),
"--seed", String(seed ?? -1),
"--model", resolvedModelPath,
];
// Add clip-on-cpu flag if enabled (required for SDXL on Metal)
if (clipOnCpu === "true") {
args.push("--clip-on-cpu");
}
// Add VAE tiling if enabled
if (vaeTiling === "true") {
args.push("--vae-tiling");
}
// Add upscaling if enabled (always run model once at native 4x, resize after)
const isUpscaled = enableUpscale === "true" && upscaleModelPath;
if (isUpscaled) {
const resolvedUpscaleModelPath = expandTilde(upscaleModelPath);
args.push("--upscale-model", resolvedUpscaleModelPath);
args.push("--upscale-repeats", "1");
args.push("--upscale-tile-size", String(upscaleTileSize));
console.log("[create-image] Upscaling enabled with model:", resolvedUpscaleModelPath, "target factor:", upscaleFactor + "x");
}
// Debug logging - will appear in lms dev console
console.log("[create-image] Executing sd-cli with args:");
console.log("[create-image] Command:", resolvedSdCliPath);
console.log("[create-image] Prompt:", prompt);
console.log("[create-image] Full args:", JSON.stringify(args, null, 2));
try {
const startTime = Date.now();
const result = await runCommand(resolvedSdCliPath, args);
const duration = ((Date.now() - startTime) / 1000).toFixed(1);
// Log sd-cli output for debugging
console.log("[create-image] sd-cli stdout:", result.stdout);
if (result.stderr) console.log("[create-image] sd-cli stderr:", result.stderr);
// Check if file was created
try {
await fs.access(tempOutputPath);
} catch {
return `Error: Image file was not created.\n\nsd-cli output:\n${result.stderr}\n${result.stdout}`;
}
// If upscaled, check actual output dimensions and resize to target factor if needed
const factor = Number(upscaleFactor);
let finalWidth = width;
let finalHeight = height;
if (isUpscaled) {
// Check actual image dimensions to verify upscale worked
try {
const sipsInfo = execSync(`sips -g pixelWidth -g pixelHeight "${tempOutputPath}"`, { stdio: 'pipe' }).toString();
const actualWidth = parseInt(sipsInfo.match(/pixelWidth:\s*(\d+)/)?.[1] || "0");
const actualHeight = parseInt(sipsInfo.match(/pixelHeight:\s*(\d+)/)?.[1] || "0");
console.log(`[create-image] Actual output dimensions: ${actualWidth}x${actualHeight}`);
if (actualWidth > width && actualHeight > height) {
// Upscale succeeded - resize to target factor if needed
finalWidth = width * factor;
finalHeight = height * factor;
if (actualWidth !== finalWidth || actualHeight !== finalHeight) {
console.log(`[create-image] Resizing from ${actualWidth}x${actualHeight} to ${finalWidth}x${finalHeight} (${factor}x)`);
try {
execSync(`python3 -c "from PIL import Image; img = Image.open('${tempOutputPath}'); img = img.resize((${finalWidth}, ${finalHeight}), Image.LANCZOS); img.save('${tempOutputPath}')"`, {
stdio: 'pipe'
});
} catch {
execSync(`sips --resampleHeightWidth ${finalHeight} ${finalWidth} "${tempOutputPath}"`, {
stdio: 'pipe'
});
}
} else {
finalWidth = actualWidth;
finalHeight = actualHeight;
}
} else {
console.log("[create-image] WARNING: Upscale failed, output is base resolution");
}
} catch (dimErr: any) {
console.error("[create-image] Failed to check dimensions:", dimErr.message);
}
}
// Read the generated image
const imageBuffer = await fs.readFile(tempOutputPath);
const fileSizeBytes = imageBuffer.length;
const fileSizeMB = (fileSizeBytes / (1024 * 1024)).toFixed(2);
console.log("[create-image] Image size:", fileSizeMB, "MB");
// Write resized preview image to LM Studio's working directory for inline display
const workingDir = ctl.getWorkingDirectory();
const previewScale = config.get("previewScale") as number;
const imageFileName = `image-${timestamp}.png`;
const workingDirImagePath = join(workingDir, imageFileName);
console.log("[create-image] Working directory:", workingDir);
console.log("[create-image] Preview scale:", previewScale, "%");
console.log("[create-image] Writing preview to:", workingDirImagePath);
try {
// Calculate preview dimensions based on final image size
const previewWidth = Math.round(finalWidth * (previewScale / 100));
const previewHeight = Math.round(finalHeight * (previewScale / 100));
console.log("[create-image] Resizing to:", previewWidth, "x", previewHeight);
// First copy the file to working directory
await fs.copyFile(tempOutputPath, workingDirImagePath);
// Use macOS sips to resize (built-in, no external deps)
execSync(`sips --resampleHeightWidth ${previewHeight} ${previewWidth} "${workingDirImagePath}"`, {
stdio: 'pipe'
});
console.log("[create-image] Preview image created successfully");
} catch (resizeError: any) {
console.error("[create-image] Failed to create preview:", resizeError.message);
// Fall back to copying full-size image
try {
await fs.copyFile(tempOutputPath, workingDirImagePath);
console.log("[create-image] Fell back to copying full-size image");
} catch (copyError: any) {
return `Image generated but could not create preview: ${resizeError.message}\n\nFull-size saved to: ${tempOutputPath}`;
}
}
// Build settings info including upscale status
const upscaleInfo = isUpscaled
? `\n🔍 Upscaled: Yes (${factor}x → ${finalWidth}x${finalHeight})`
: "";
const settingsInfo = `Prompt: ${prompt}\nSize: ${width}x${height}${upscaleInfo}\nSteps: ${steps}\nSampler: ${sampler}\nScheduler: ${scheduler}\nCFG Scale: ${cfgScale}\nSeed: ${seed ?? "random"}`;
// Return MCP-style format with markdown preview and link to full-size
return [
{
type: "image",
fileName: imageFileName,
mimeType: "image/png",
markdown: ``,
"$hint": "This is an image file. Present the image to the user by using the markdown above."
},
{
type: "text",
text: `✨ Image generated successfully! (${duration}s)\n\n📋 Settings:\n${settingsInfo}\n\n🖼️ Preview: ${previewScale}% scale\n📁 Full-size: [${tempOutputPath}](file://${tempOutputPath})\n📦 Size: ${fileSizeMB}MB`
}
];
} catch (error: any) {
console.error("[create-image] Error:", error);
return `❌ Error generating image: ${error.message}\n\nDetails:\n${error.stderr || error.stdout || "No additional details"}`;
}
},
});
tools.push(generateImageTool);
return tools;
}
function runCommand(command: string, args: string[]): Promise<{ stdout: string; stderr: string }> {
return new Promise((resolve, reject) => {
const proc = spawn(command, args);
let stdout = "";
let stderr = "";
proc.stdout.on("data", (data) => {
stdout += data.toString();
});
proc.stderr.on("data", (data) => {
stderr += data.toString();
});
proc.on("close", (code, signal) => {
// CLI has a cleanup bug: crashes with SIGABRT in ggml_metal_rsets_free after successful save
// code=134 or code=null+signal=SIGABRT — treat both as success
if (code === 0 || code === 134 || signal === "SIGABRT") {
resolve({ stdout, stderr });
} else {
reject({ message: `Command failed with code ${code}`, stdout, stderr });
}
});
proc.on("error", (error) => {
reject({ message: error.message, stdout, stderr });
});
});
}
src / toolsProvider.ts
import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./config";
import { spawn, execSync } from "child_process";
import { promises as fs } from "fs";
import { join } from "path";
import { homedir } from "os";
function expandTilde(filepath: string): string {
if (filepath.startsWith("~/") || filepath === "~") {
return join(homedir(), filepath.slice(1));
}
return filepath;
}
export async function toolsProvider(ctl: ToolsProviderController) {
const config = ctl.getPluginConfig(configSchematics);
const tools: Tool[] = [];
const generateImageTool = tool({
name: "generate_image",
description: "Generate an image using Stable Diffusion from a text prompt. All parameters except 'prompt' are optional and have pre-configured defaults. DO NOT set optional parameters unless the user explicitly specifies them in their message. Omitting a parameter uses the user's configured default.",
parameters: {
prompt: z.string().describe("Text description of the image to generate"),
negative_prompt: z.string().optional().describe("What to avoid in the image. Omit to use default."),
seed: z.number().optional().describe("Random seed. Omit for random."),
width: z.number().optional().describe("Image width in pixels (multiple of 64, 256-1536). OMIT unless user specifies."),
height: z.number().optional().describe("Image height in pixels (multiple of 64, 256-1536). OMIT unless user specifies."),
steps: z.number().optional().describe("Sampling steps (1-50). OMIT unless user specifies."),
cfg_scale: z.number().optional().describe("CFG scale (1-20). OMIT unless user specifies."),
sampler: z.enum(["lcm", "euler_a", "euler", "dpm++2m", "dpm++2s_a", "heun"]).optional().describe("Sampling method. OMIT unless user specifies."),
scheduler: z.enum(["discrete", "karras", "exponential", "ays", "gits", "smoothstep", "sgm_uniform", "simple", "kl_optimal", "lcm", "bong_tangent"]).optional().describe("Scheduler. OMIT unless user specifies."),
upscale: z.union([z.boolean(), z.string()]).optional().describe("Whether to upscale. OMIT unless user specifies."),
upscale_factor: z.number().optional().describe("Upscale factor: 2, 3, or 4. OMIT unless user specifies.")
},
implementation: async ({ prompt, negative_prompt, seed, width: paramWidth, height: paramHeight, steps: paramSteps, cfg_scale: paramCfgScale, sampler: paramSampler, scheduler: paramScheduler, upscale: paramUpscale, upscale_factor: paramUpscaleFactorNum }) => {
// Read config values fresh on each invocation so UI changes take effect
// LLM-provided parameters override config defaults
const sdCliPath = config.get("sdCliPath");
const modelPath = config.get("modelPath");
const outputDirectory = config.get("outputDirectory");
const width = paramWidth ?? config.get("width");
const height = paramHeight ?? config.get("height");
const steps = paramSteps ?? config.get("steps");
const cfgScale = paramCfgScale ?? config.get("cfgScale");
const sampler = paramSampler ?? config.get("sampler");
const scheduler = paramScheduler ?? config.get("scheduler");
const clipOnCpu = config.get("clipOnCpu");
const vaeTiling = config.get("vaeTiling");
const enableUpscale = paramUpscale !== undefined ? (String(paramUpscale).toLowerCase() === "true" ? "true" : "false") : config.get("enableUpscale");
const upscaleModelPath = config.get("upscaleModelPath");
const upscaleTileSize = config.get("upscaleTileSize");
const upscaleFactor = paramUpscaleFactorNum !== undefined ? String(paramUpscaleFactorNum) : config.get("upscaleFactor");
// Validate configuration
if (!sdCliPath) {
return "Error: sd-cli path not configured. Please set it in the plugin settings (click on create-image in Integrations panel).";
}
if (!modelPath) {
return "Error: Model path not configured. Please set it in the plugin settings.";
}
// Resolve paths (expand ~ to home directory)
const resolvedSdCliPath = expandTilde(sdCliPath);
const resolvedModelPath = expandTilde(modelPath);
// Generate unique filename for sd-cli output
const timestamp = Date.now();
const resolvedOutputDir = expandTilde(outputDirectory || "/tmp");
const tempOutputPath = join(resolvedOutputDir, `sd_${timestamp}.png`);
// Ensure output directory exists
await fs.mkdir(resolvedOutputDir, { recursive: true });
// Default negative prompt
const negPrompt = negative_prompt || "blurry, bad quality, distorted, ugly, low resolution, watermark, text";
// Build command arguments
const args = [
"--prompt", prompt,
"--negative-prompt", negPrompt,
"--output", tempOutputPath,
"--width", String(width),
"--height", String(height),
"--steps", String(steps),
"--sampling-method", sampler,
"--scheduler", scheduler,
"--cfg-scale", String(cfgScale),
"--seed", String(seed ?? -1),
"--model", resolvedModelPath,
];
// Add clip-on-cpu flag if enabled (required for SDXL on Metal)
if (clipOnCpu === "true") {
args.push("--clip-on-cpu");
}
// Add VAE tiling if enabled
if (vaeTiling === "true") {
args.push("--vae-tiling");
}
// Add upscaling if enabled (always run model once at native 4x, resize after)
const isUpscaled = enableUpscale === "true" && upscaleModelPath;
if (isUpscaled) {
const resolvedUpscaleModelPath = expandTilde(upscaleModelPath);
args.push("--upscale-model", resolvedUpscaleModelPath);
args.push("--upscale-repeats", "1");
args.push("--upscale-tile-size", String(upscaleTileSize));
console.log("[create-image] Upscaling enabled with model:", resolvedUpscaleModelPath, "target factor:", upscaleFactor + "x");
}
// Debug logging - will appear in lms dev console
console.log("[create-image] Executing sd-cli with args:");
console.log("[create-image] Command:", resolvedSdCliPath);
console.log("[create-image] Prompt:", prompt);
console.log("[create-image] Full args:", JSON.stringify(args, null, 2));
try {
const startTime = Date.now();
const result = await runCommand(resolvedSdCliPath, args);
const duration = ((Date.now() - startTime) / 1000).toFixed(1);
// Log sd-cli output for debugging
console.log("[create-image] sd-cli stdout:", result.stdout);
if (result.stderr) console.log("[create-image] sd-cli stderr:", result.stderr);
// Check if file was created
try {
await fs.access(tempOutputPath);
} catch {
return `Error: Image file was not created.\n\nsd-cli output:\n${result.stderr}\n${result.stdout}`;
}
// If upscaled, check actual output dimensions and resize to target factor if needed
const factor = Number(upscaleFactor);
let finalWidth = width;
let finalHeight = height;
if (isUpscaled) {
// Check actual image dimensions to verify upscale worked
try {
const sipsInfo = execSync(`sips -g pixelWidth -g pixelHeight "${tempOutputPath}"`, { stdio: 'pipe' }).toString();
const actualWidth = parseInt(sipsInfo.match(/pixelWidth:\s*(\d+)/)?.[1] || "0");
const actualHeight = parseInt(sipsInfo.match(/pixelHeight:\s*(\d+)/)?.[1] || "0");
console.log(`[create-image] Actual output dimensions: ${actualWidth}x${actualHeight}`);
if (actualWidth > width && actualHeight > height) {
// Upscale succeeded - resize to target factor if needed
finalWidth = width * factor;
finalHeight = height * factor;
if (actualWidth !== finalWidth || actualHeight !== finalHeight) {
console.log(`[create-image] Resizing from ${actualWidth}x${actualHeight} to ${finalWidth}x${finalHeight} (${factor}x)`);
try {
execSync(`python3 -c "from PIL import Image; img = Image.open('${tempOutputPath}'); img = img.resize((${finalWidth}, ${finalHeight}), Image.LANCZOS); img.save('${tempOutputPath}')"`, {
stdio: 'pipe'
});
} catch {
execSync(`sips --resampleHeightWidth ${finalHeight} ${finalWidth} "${tempOutputPath}"`, {
stdio: 'pipe'
});
}
} else {
finalWidth = actualWidth;
finalHeight = actualHeight;
}
} else {
console.log("[create-image] WARNING: Upscale failed, output is base resolution");
}
} catch (dimErr: any) {
console.error("[create-image] Failed to check dimensions:", dimErr.message);
}
}
// Read the generated image
const imageBuffer = await fs.readFile(tempOutputPath);
const fileSizeBytes = imageBuffer.length;
const fileSizeMB = (fileSizeBytes / (1024 * 1024)).toFixed(2);
console.log("[create-image] Image size:", fileSizeMB, "MB");
// Write resized preview image to LM Studio's working directory for inline display
const workingDir = ctl.getWorkingDirectory();
const previewScale = config.get("previewScale") as number;
const imageFileName = `image-${timestamp}.png`;
const workingDirImagePath = join(workingDir, imageFileName);
console.log("[create-image] Working directory:", workingDir);
console.log("[create-image] Preview scale:", previewScale, "%");
console.log("[create-image] Writing preview to:", workingDirImagePath);
try {
// Calculate preview dimensions based on final image size
const previewWidth = Math.round(finalWidth * (previewScale / 100));
const previewHeight = Math.round(finalHeight * (previewScale / 100));
console.log("[create-image] Resizing to:", previewWidth, "x", previewHeight);
// First copy the file to working directory
await fs.copyFile(tempOutputPath, workingDirImagePath);
// Use macOS sips to resize (built-in, no external deps)
execSync(`sips --resampleHeightWidth ${previewHeight} ${previewWidth} "${workingDirImagePath}"`, {
stdio: 'pipe'
});
console.log("[create-image] Preview image created successfully");
} catch (resizeError: any) {
console.error("[create-image] Failed to create preview:", resizeError.message);
// Fall back to copying full-size image
try {
await fs.copyFile(tempOutputPath, workingDirImagePath);
console.log("[create-image] Fell back to copying full-size image");
} catch (copyError: any) {
return `Image generated but could not create preview: ${resizeError.message}\n\nFull-size saved to: ${tempOutputPath}`;
}
}
// Build settings info including upscale status
const upscaleInfo = isUpscaled
? `\n🔍 Upscaled: Yes (${factor}x → ${finalWidth}x${finalHeight})`
: "";
const settingsInfo = `Prompt: ${prompt}\nSize: ${width}x${height}${upscaleInfo}\nSteps: ${steps}\nSampler: ${sampler}\nScheduler: ${scheduler}\nCFG Scale: ${cfgScale}\nSeed: ${seed ?? "random"}`;
// Return MCP-style format with markdown preview and link to full-size
return [
{
type: "image",
fileName: imageFileName,
mimeType: "image/png",
markdown: ``,
"$hint": "This is an image file. Present the image to the user by using the markdown above."
},
{
type: "text",
text: `✨ Image generated successfully! (${duration}s)\n\n📋 Settings:\n${settingsInfo}\n\n🖼️ Preview: ${previewScale}% scale\n📁 Full-size: [${tempOutputPath}](file://${tempOutputPath})\n📦 Size: ${fileSizeMB}MB`
}
];
} catch (error: any) {
console.error("[create-image] Error:", error);
return `❌ Error generating image: ${error.message}\n\nDetails:\n${error.stderr || error.stdout || "No additional details"}`;
}
},
});
tools.push(generateImageTool);
return tools;
}
function runCommand(command: string, args: string[]): Promise<{ stdout: string; stderr: string }> {
return new Promise((resolve, reject) => {
const proc = spawn(command, args);
let stdout = "";
let stderr = "";
proc.stdout.on("data", (data) => {
stdout += data.toString();
});
proc.stderr.on("data", (data) => {
stderr += data.toString();
});
proc.on("close", (code, signal) => {
// CLI has a cleanup bug: crashes with SIGABRT in ggml_metal_rsets_free after successful save
// code=134 or code=null+signal=SIGABRT — treat both as success
if (code === 0 || code === 134 || signal === "SIGABRT") {
resolve({ stdout, stderr });
} else {
reject({ message: `Command failed with code ${code}`, stdout, stderr });
}
});
proc.on("error", (error) => {
reject({ message: error.message, stdout, stderr });
});
});
}