genai-electron 0.3.0

package/dist/managers/DiffusionServerManager.js

@@ -0,0 +1,649 @@
+import { ServerManager } from './ServerManager.js';
+import { ModelManager } from './ModelManager.js';
+import { SystemInfo } from '../system/SystemInfo.js';
+import { ProcessManager } from '../process/ProcessManager.js';
+import { ResourceOrchestrator } from './ResourceOrchestrator.js';
+import { GenerationRegistry } from './GenerationRegistry.js';
+import http from 'node:http';
+import { promises as fs } from 'node:fs';
+import { getTempPath } from '../config/paths.js';
+import { BINARY_VERSIONS, DEFAULT_PORTS } from '../config/defaults.js';
+import { deleteFile } from '../utils/file-utils.js';
+import { ServerError, ModelNotFoundError, InsufficientResourcesError } from '../errors/index.js';
+export class DiffusionServerManager extends ServerManager {
+    processManager;
+    modelManager;
+    systemInfo;
+    orchestrator;
+    registry;
+    binaryPath;
+    httpServer;
+    currentGeneration;
+    currentModelInfo;
+    modelLoadTime = 2000;
+    diffusionTimePerStepPerMegapixel = 1000;
+    vaeTimePerMegapixel = 8000;
+    generationStartTime;
+    loadStartTime;
+    loadEndTime;
+    diffusionStartTime;
+    diffusionEndTime;
+    vaeStartTime;
+    vaeEndTime;
+    syntheticProgressInterval;
+    currentStage;
+    totalEstimatedTime = 0;
+    loadProgress = { current: 0, total: 0 };
+    diffusionProgress = { current: 0, total: 0 };
+    constructor(modelManager = ModelManager.getInstance(), systemInfo = SystemInfo.getInstance(), llamaServer) {
+        super();
+        this.processManager = new ProcessManager();
+        this.modelManager = modelManager;
+        this.systemInfo = systemInfo;
+        this.registry = new GenerationRegistry({
+            maxResultAgeMs: parseInt(process.env.IMAGE_RESULT_TTL_MS || '300000', 10),
+            cleanupIntervalMs: parseInt(process.env.IMAGE_CLEANUP_INTERVAL_MS || '60000', 10),
+        });
+        if (llamaServer) {
+            this.orchestrator = new ResourceOrchestrator(systemInfo, llamaServer, this, modelManager);
+        }
+    }
+    async start(config) {
+        if (this._status === 'running') {
+            throw new ServerError('Server is already running', {
+                suggestion: 'Stop the server first with stop()',
+            });
+        }
+        this.setStatus('starting');
+        this._config = config;
+        try {
+            const modelInfo = await this.modelManager.getModelInfo(config.modelId);
+            if (modelInfo.type !== 'diffusion') {
+                throw new ModelNotFoundError(`Model ${config.modelId} is not a diffusion model (type: ${modelInfo.type})`);
+            }
+            this.currentModelInfo = modelInfo;
+            const canRun = await this.systemInfo.canRunModel(modelInfo, { checkTotalMemory: true });
+            if (!canRun.possible) {
+                const memoryInfo = this.systemInfo.getMemoryInfo();
+                throw new InsufficientResourcesError(`System cannot run model: ${canRun.reason || 'Insufficient resources'}`, {
+                    required: `Model size: ${Math.round(modelInfo.size / 1024 / 1024 / 1024)}GB`,
+                    available: `Total RAM: ${Math.round(memoryInfo.total / 1024 / 1024 / 1024)}GB`,
+                    suggestion: canRun.suggestion || canRun.reason || 'Try a smaller model',
+                });
+            }
+            this.binaryPath = await this.ensureBinary(modelInfo.path, config.forceValidation);
+            const port = config.port || DEFAULT_PORTS.diffusion;
+            await this.checkPortAvailability(port);
+            await this.initializeLogManager('diffusion-server.log', `Starting diffusion server on port ${port}`);
+            await this.createHTTPServer(config);
+            this._port = port;
+            this._startedAt = new Date();
+            this.setStatus('running');
+            await this.logManager.write('Diffusion server is running', 'info');
+            this.systemInfo.clearCache();
+            this.emitEvent('started', this.getInfo());
+            return this.getInfo();
+        }
+        catch (error) {
+            throw await this.handleStartupError('diffusion-server', error, async () => {
+                if (this.httpServer) {
+                    this.httpServer.close();
+                    this.httpServer = undefined;
+                }
+            });
+        }
+    }
+    async stop() {
+        if (this._status === 'stopped') {
+            return;
+        }
+        this.setStatus('stopping');
+        try {
+            if (this.logManager) {
+                await this.logManager.write('Stopping diffusion server...', 'info');
+            }
+            if (this.currentGeneration) {
+                this.currentGeneration.cancel();
+                this.currentGeneration = undefined;
+            }
+            if (this.httpServer) {
+                await new Promise((resolve) => {
+                    this.httpServer.close(() => resolve());
+                });
+                this.httpServer = undefined;
+            }
+            this.registry.destroy();
+            this.setStatus('stopped');
+            this._port = 0;
+            if (this.logManager) {
+                await this.logManager.write('Diffusion server stopped', 'info');
+            }
+            this.systemInfo.clearCache();
+            this.emitEvent('stopped');
+        }
+        catch (error) {
+            this.setStatus('stopped');
+            throw new ServerError(`Failed to stop server: ${error instanceof Error ? error.message : 'Unknown error'}`, { error: error instanceof Error ? error.message : String(error) });
+        }
+    }
+    async generateImage(config) {
+        if (this._status !== 'running') {
+            throw new ServerError('Server is not running', {
+                suggestion: 'Start the server first with start()',
+            });
+        }
+        if (this.currentGeneration) {
+            throw new ServerError('Server is busy generating another image', {
+                suggestion: 'Wait for current generation to complete',
+            });
+        }
+        console.log('[DiffusionServer] generateImage called');
+        console.log('[DiffusionServer] orchestrator exists:', !!this.orchestrator);
+        if (this.orchestrator) {
+            console.log('[DiffusionServer] Using orchestrator for automatic resource management');
+            return this.orchestrator.orchestrateImageGeneration(config);
+        }
+        else {
+            console.log('[DiffusionServer] No orchestrator - using direct execution');
+            return this.executeImageGeneration(config);
+        }
+    }
+    async isHealthy() {
+        return this._status === 'running' && this.httpServer !== undefined;
+    }
+    getInfo() {
+        const baseInfo = super.getInfo();
+        return {
+            ...baseInfo,
+            busy: !!this.currentGeneration,
+        };
+    }
+    async ensureBinary(modelPath, forceValidation = false) {
+        return this.ensureBinaryHelper('diffusion', 'sd', BINARY_VERSIONS.diffusionCpp, modelPath, forceValidation);
+    }
+    async createHTTPServer(config) {
+        const port = config.port || DEFAULT_PORTS.diffusion;
+        this.httpServer = http.createServer(async (req, res) => {
+            res.setHeader('Access-Control-Allow-Origin', '*');
+            res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS, DELETE');
+            res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
+            if (req.method === 'OPTIONS') {
+                res.writeHead(200);
+                res.end();
+                return;
+            }
+            try {
+                if (req.url === '/health' && req.method === 'GET') {
+                    res.writeHead(200, { 'Content-Type': 'application/json' });
+                    res.end(JSON.stringify({ status: 'ok', busy: !!this.currentGeneration }));
+                    return;
+                }
+                if (req.url === '/v1/images/generations' && req.method === 'POST') {
+                    await this.handleStartGeneration(req, res);
+                    return;
+                }
+                const getMatch = req.url?.match(/^\/v1\/images\/generations\/([^/]+)$/);
+                if (getMatch && getMatch[1] && req.method === 'GET') {
+                    const generationId = getMatch[1];
+                    await this.handleGetGeneration(generationId, res);
+                    return;
+                }
+                res.writeHead(404, { 'Content-Type': 'application/json' });
+                res.end(JSON.stringify({ error: { message: 'Not found', code: 'NOT_FOUND' } }));
+            }
+            catch (error) {
+                res.writeHead(500, { 'Content-Type': 'application/json' });
+                res.end(JSON.stringify({
+                    error: {
+                        message: error instanceof Error ? error.message : 'Internal server error',
+                        code: 'INTERNAL_ERROR',
+                    },
+                }));
+            }
+        });
+        await new Promise((resolve, reject) => {
+            this.httpServer.listen(port, () => resolve());
+            this.httpServer.on('error', reject);
+        });
+        await this.logManager?.write(`HTTP server listening on port ${port}`, 'info');
+    }
+    async handleStartGeneration(req, res) {
+        const body = await this.parseRequestBody(req);
+        const imageConfig = JSON.parse(body);
+        if (!imageConfig.prompt) {
+            res.writeHead(400, { 'Content-Type': 'application/json' });
+            res.end(JSON.stringify({
+                error: { message: 'Missing required field: prompt', code: 'INVALID_REQUEST' },
+            }));
+            return;
+        }
+        if (imageConfig.count !== undefined) {
+            if (imageConfig.count < 1 || imageConfig.count > 5) {
+                res.writeHead(400, { 'Content-Type': 'application/json' });
+                res.end(JSON.stringify({
+                    error: { message: 'count must be between 1 and 5', code: 'INVALID_REQUEST' },
+                }));
+                return;
+            }
+        }
+        if (this.currentGeneration) {
+            res.writeHead(503, { 'Content-Type': 'application/json' });
+            res.end(JSON.stringify({
+                error: {
+                    message: 'Server is busy generating another image',
+                    code: 'SERVER_BUSY',
+                    suggestion: 'Wait for current generation to complete and try again',
+                },
+            }));
+            return;
+        }
+        const id = this.registry.create(imageConfig);
+        this.runAsyncGeneration(id, imageConfig).catch((error) => {
+            this.registry.update(id, {
+                status: 'error',
+                error: {
+                    message: error instanceof Error ? error.message : 'Unknown error',
+                    code: this.mapErrorCode(error),
+                },
+            });
+        });
+        res.writeHead(201, { 'Content-Type': 'application/json' });
+        res.end(JSON.stringify({
+            id,
+            status: 'pending',
+            createdAt: Date.now(),
+        }));
+    }
+    async handleGetGeneration(id, res) {
+        const state = this.registry.get(id);
+        if (!state) {
+            res.writeHead(404, { 'Content-Type': 'application/json' });
+            res.end(JSON.stringify({ error: { message: 'Generation not found', code: 'NOT_FOUND' } }));
+            return;
+        }
+        const response = {
+            id: state.id,
+            status: state.status,
+            createdAt: state.createdAt,
+            updatedAt: state.updatedAt,
+        };
+        if (state.status === 'in_progress' && state.progress) {
+            response.progress = state.progress;
+        }
+        if (state.status === 'complete' && state.result) {
+            response.result = state.result;
+        }
+        if (state.status === 'error' && state.error) {
+            response.error = state.error;
+        }
+        res.writeHead(200, { 'Content-Type': 'application/json' });
+        res.end(JSON.stringify(response));
+    }
+    async runAsyncGeneration(id, config) {
+        const startTime = Date.now();
+        this.registry.update(id, { status: 'in_progress' });
+        const wrappedConfig = {
+            ...config,
+            onProgress: (currentStep, totalSteps, stage, percentage) => {
+                this.registry.update(id, {
+                    progress: {
+                        currentStep,
+                        totalSteps,
+                        stage,
+                        percentage,
+                        currentImage: config.count && config.count > 1
+                            ? Math.floor((percentage || 0) / (100 / config.count)) + 1
+                            : undefined,
+                        totalImages: config.count && config.count > 1 ? config.count : undefined,
+                    },
+                });
+                config.onProgress?.(currentStep, totalSteps, stage, percentage);
+            },
+        };
+        const count = config.count || 1;
+        let results;
+        if (count > 1) {
+            results = await this.executeBatchGeneration(wrappedConfig);
+        }
+        else {
+            if (this.orchestrator) {
+                results = [await this.orchestrator.orchestrateImageGeneration(wrappedConfig)];
+            }
+            else {
+                results = [await this.executeImageGeneration(wrappedConfig)];
+            }
+        }
+        const images = results.map((result) => ({
+            image: result.image.toString('base64'),
+            seed: result.seed,
+            width: result.width,
+            height: result.height,
+        }));
+        this.registry.update(id, {
+            status: 'complete',
+            result: {
+                images,
+                format: 'png',
+                timeTaken: Date.now() - startTime,
+            },
+        });
+    }
+    mapErrorCode(error) {
+        if (error instanceof Error) {
+            const message = error.message.toLowerCase();
+            if (message.includes('server is busy'))
+                return 'SERVER_BUSY';
+            if (message.includes('not running'))
+                return 'SERVER_NOT_RUNNING';
+            if (message.includes('failed to spawn'))
+                return 'BACKEND_ERROR';
+            if (message.includes('exited with code'))
+                return 'BACKEND_ERROR';
+            if (message.includes('failed to read'))
+                return 'IO_ERROR';
+        }
+        return 'UNKNOWN_ERROR';
+    }
+    parseRequestBody(req) {
+        return new Promise((resolve, reject) => {
+            let body = '';
+            req.on('data', (chunk) => {
+                body += chunk.toString();
+            });
+            req.on('end', () => resolve(body));
+            req.on('error', reject);
+        });
+    }
+    async executeImageGeneration(config) {
+        const startTime = Date.now();
+        if (!this.currentModelInfo) {
+            throw new ServerError('Model information not available', {
+                suggestion: 'This is an internal error - model should have been loaded',
+            });
+        }
+        const normalizedConfig = {
+            ...config,
+            seed: config.seed === undefined || config.seed < 0 ? this.generateRandomSeed() : config.seed,
+        };
+        this.initializeProgressTracking(normalizedConfig);
+        const args = this.buildDiffusionArgs(normalizedConfig, this.currentModelInfo);
+        const outputPath = getTempPath(`sd-output-${Date.now()}.png`);
+        args.push('-o', outputPath);
+        await this.logManager?.write(`Generating image: ${this.binaryPath} ${args.join(' ')}`, 'info');
+        let cancelled = false;
+        let pid;
+        const generationPromise = new Promise((resolve, reject) => {
+            const spawnResult = this.processManager.spawn(this.binaryPath, args, {
+                onStdout: (data) => {
+                    this.processStdoutForProgress(data, normalizedConfig);
+                    this.logManager?.write(data, 'info').catch(() => void 0);
+                },
+                onStderr: (data) => {
+                    this.logManager?.write(data, 'warn').catch(() => void 0);
+                },
+                onExit: async (code) => {
+                    this.cleanupSyntheticProgress();
+                    if (cancelled) {
+                        reject(new Error('Image generation cancelled'));
+                        return;
+                    }
+                    if (code !== 0) {
+                        reject(new ServerError(`stable-diffusion.cpp exited with code ${code}`));
+                        return;
+                    }
+                    try {
+                        const imageBuffer = await fs.readFile(outputPath);
+                        await deleteFile(outputPath).catch(() => void 0);
+                        this.updateTimeEstimates(normalizedConfig);
+                        resolve({
+                            image: imageBuffer,
+                            format: 'png',
+                            timeTaken: Date.now() - startTime,
+                            seed: normalizedConfig.seed,
+                            width: normalizedConfig.width || 512,
+                            height: normalizedConfig.height || 512,
+                        });
+                    }
+                    catch (error) {
+                        reject(new ServerError('Failed to read generated image', {
+                            error: error instanceof Error ? error.message : String(error),
+                        }));
+                    }
+                },
+                onError: (error) => {
+                    this.cleanupSyntheticProgress();
+                    reject(new ServerError('Failed to spawn stable-diffusion.cpp', {
+                        error: error.message,
+                    }));
+                },
+            });
+            pid = spawnResult.pid;
+        });
+        this.currentGeneration = {
+            promise: generationPromise,
+            cancel: () => {
+                cancelled = true;
+                this.cleanupSyntheticProgress();
+                if (pid !== undefined) {
+                    this.processManager.kill(pid, 5000).catch(() => void 0);
+                }
+            },
+        };
+        try {
+            const result = await generationPromise;
+            this.currentGeneration = undefined;
+            return result;
+        }
+        catch (error) {
+            this.currentGeneration = undefined;
+            throw error;
+        }
+    }
+    async executeBatchGeneration(config) {
+        const count = config.count || 1;
+        const images = [];
+        for (let i = 0; i < count; i++) {
+            const imageSeed = config.seed !== undefined && config.seed >= 0 ? config.seed + i : this.generateRandomSeed();
+            const wrappedConfig = {
+                ...config,
+                seed: imageSeed,
+                onProgress: config.onProgress
+                    ? (currentStep, totalSteps, stage, percentage) => {
+                        const completedImages = i;
+                        const currentImageProgress = (percentage || 0) / 100;
+                        const overallPercentage = ((completedImages + currentImageProgress) / count) * 100;
+                        config.onProgress(currentStep, totalSteps, stage, overallPercentage);
+                    }
+                    : undefined,
+            };
+            const result = await this.executeImageGeneration(wrappedConfig);
+            images.push(result);
+        }
+        return images;
+    }
+    buildDiffusionArgs(config, modelInfo) {
+        const args = [];
+        args.push('-m', modelInfo.path);
+        if (config.prompt) {
+            args.push('-p', config.prompt);
+        }
+        if (config.negativePrompt) {
+            args.push('-n', config.negativePrompt);
+        }
+        if (config.width) {
+            args.push('-W', String(config.width));
+        }
+        if (config.height) {
+            args.push('-H', String(config.height));
+        }
+        if (config.steps) {
+            args.push('--steps', String(config.steps));
+        }
+        if (config.cfgScale) {
+            args.push('--cfg-scale', String(config.cfgScale));
+        }
+        if (config.seed !== undefined) {
+            args.push('-s', String(config.seed));
+        }
+        if (config.sampler) {
+            args.push('--sampling-method', config.sampler);
+        }
+        const serverConfig = this._config;
+        if (serverConfig.gpuLayers !== undefined && serverConfig.gpuLayers > 0) {
+            args.push('--n-gpu-layers', String(serverConfig.gpuLayers));
+        }
+        if (serverConfig.threads) {
+            args.push('-t', String(serverConfig.threads));
+        }
+        return args;
+    }
+    generateRandomSeed() {
+        return Math.floor(Math.random() * 2147483647);
+    }
+    initializeProgressTracking(config) {
+        const width = config.width || 512;
+        const height = config.height || 512;
+        const steps = config.steps || 20;
+        const megapixels = (width * height) / 1_000_000;
+        this.totalEstimatedTime =
+            this.modelLoadTime +
+                steps * megapixels * this.diffusionTimePerStepPerMegapixel +
+                megapixels * this.vaeTimePerMegapixel;
+        this.generationStartTime = Date.now();
+        this.currentStage = undefined;
+        this.loadStartTime = undefined;
+        this.loadEndTime = undefined;
+        this.diffusionStartTime = undefined;
+        this.diffusionEndTime = undefined;
+        this.vaeStartTime = undefined;
+        this.vaeEndTime = undefined;
+        this.loadProgress = { current: 0, total: 0 };
+        this.diffusionProgress = { current: 0, total: 0 };
+    }
+    processStdoutForProgress(data, config) {
+        if (data.includes('loading tensors from')) {
+            this.currentStage = 'loading';
+            this.loadStartTime = Date.now();
+        }
+        else if (data.includes('generating image:') || data.includes('sampling using')) {
+            if (this.currentStage === 'loading') {
+                this.loadEndTime = Date.now();
+            }
+            this.currentStage = 'diffusion';
+            this.diffusionStartTime = Date.now();
+        }
+        else if (data.includes('decoding 1 latents')) {
+            if (this.currentStage === 'diffusion') {
+                this.diffusionEndTime = Date.now();
+            }
+            this.currentStage = 'vae';
+            this.vaeStartTime = Date.now();
+            this.startSyntheticVaeProgress(config);
+        }
+        else if (data.includes('decode_first_stage completed')) {
+            this.vaeEndTime = Date.now();
+            this.cleanupSyntheticProgress();
+            if (config.onProgress) {
+                config.onProgress(0, 0, 'decoding', 100);
+            }
+        }
+        const progressMatch = data.match(/\|\s*(\d+)\/(\d+)\s*-/);
+        if (progressMatch && progressMatch[1] && progressMatch[2]) {
+            const current = parseInt(progressMatch[1], 10);
+            const total = parseInt(progressMatch[2], 10);
+            if (this.currentStage === 'loading') {
+                this.loadProgress = { current, total };
+                this.reportProgress(config);
+            }
+            else if (this.currentStage === 'diffusion') {
+                this.diffusionProgress = { current, total };
+                this.reportProgress(config);
+            }
+            else if (this.currentStage === undefined) {
+                this.currentStage = 'loading';
+                this.loadStartTime = Date.now();
+                this.loadProgress = { current, total };
+                this.reportProgress(config);
+            }
+        }
+    }
+    reportProgress(config) {
+        if (!config.onProgress || !this.generationStartTime)
+            return;
+        const percentage = this.calculateOverallPercentage();
+        if (this.currentStage === 'loading' && this.loadProgress.total > 0) {
+            config.onProgress(this.loadProgress.current, this.loadProgress.total, 'loading', percentage);
+        }
+        else if (this.currentStage === 'diffusion' && this.diffusionProgress.total > 0) {
+            config.onProgress(this.diffusionProgress.current, this.diffusionProgress.total, 'diffusion', percentage);
+        }
+        else if (this.currentStage === 'vae') {
+            config.onProgress(0, 0, 'decoding', percentage);
+        }
+    }
+    calculateOverallPercentage() {
+        if (!this.generationStartTime || this.totalEstimatedTime === 0)
+            return 0;
+        let elapsedTotal = 0;
+        if (this.currentStage === 'loading') {
+            const elapsedLoad = Date.now() - (this.loadStartTime || this.generationStartTime);
+            elapsedTotal = elapsedLoad;
+        }
+        else if (this.currentStage === 'diffusion') {
+            const actualLoadTime = this.loadEndTime
+                ? this.loadEndTime - (this.loadStartTime || this.generationStartTime)
+                : this.modelLoadTime;
+            const elapsedDiffusion = Date.now() - (this.diffusionStartTime || Date.now());
+            elapsedTotal = actualLoadTime + elapsedDiffusion;
+        }
+        else if (this.currentStage === 'vae') {
+            const actualLoadTime = this.loadEndTime
+                ? this.loadEndTime - (this.loadStartTime || this.generationStartTime)
+                : this.modelLoadTime;
+            const actualDiffusionTime = this.diffusionEndTime
+                ? this.diffusionEndTime - (this.diffusionStartTime || Date.now())
+                : 0;
+            const elapsedVae = Date.now() - (this.vaeStartTime || Date.now());
+            elapsedTotal = actualLoadTime + actualDiffusionTime + elapsedVae;
+        }
+        return Math.min(100, Math.round((elapsedTotal / this.totalEstimatedTime) * 100));
+    }
+    startSyntheticVaeProgress(config) {
+        if (!config.onProgress)
+            return;
+        this.cleanupSyntheticProgress();
+        this.syntheticProgressInterval = setInterval(() => {
+            const percentage = this.calculateOverallPercentage();
+            config.onProgress(0, 0, 'decoding', percentage);
+        }, 100);
+        if (this.syntheticProgressInterval &&
+            typeof this.syntheticProgressInterval.unref === 'function') {
+            this.syntheticProgressInterval.unref();
+        }
+    }
+    cleanupSyntheticProgress() {
+        if (this.syntheticProgressInterval) {
+            clearInterval(this.syntheticProgressInterval);
+            this.syntheticProgressInterval = undefined;
+        }
+    }
+    updateTimeEstimates(config) {
+        const width = config.width || 512;
+        const height = config.height || 512;
+        const steps = config.steps || 20;
+        const megapixels = (width * height) / 1_000_000;
+        if (this.loadStartTime && this.loadEndTime) {
+            const actualLoadTime = this.loadEndTime - this.loadStartTime;
+            this.modelLoadTime = actualLoadTime;
+        }
+        if (this.diffusionStartTime && this.diffusionEndTime) {
+            const actualDiffusionTime = this.diffusionEndTime - this.diffusionStartTime;
+            this.diffusionTimePerStepPerMegapixel = actualDiffusionTime / (steps * megapixels);
+        }
+        if (this.vaeStartTime && this.vaeEndTime) {
+            const actualVaeTime = this.vaeEndTime - this.vaeStartTime;
+            this.vaeTimePerMegapixel = actualVaeTime / megapixels;
+        }
+    }
+}
+//# sourceMappingURL=DiffusionServerManager.js.map