import { connectToComfyUI, queuePrompt, getImage, getHistory, disconnect } from '../src/comfyui-client.js';
import { getUpscalingWorkflow, selectBestModel } from '../src/workflows/upscaling.js';
import fs from 'fs';
import sizeOf from 'image-size';
async function testUpscaling() {
console.log('Testing upscaling functionality...\n');
try {
// Connect to ComfyUI
await connectToComfyUI('comfyui', 8188);
console.log('✅ Connected to ComfyUI\n');
// Test image path
const testImage = 'flux_output_00001_.png';
// Check original image dimensions
const originalPath = `output/${testImage}`;
if (fs.existsSync(originalPath)) {
const dimensions = sizeOf(originalPath);
console.log(`Original image: ${testImage}`);
console.log(`Dimensions: ${dimensions.width}x${dimensions.height}\n`);
}
// Test different models
const tests = [
{ model: '4x-UltraSharp.safetensors', description: 'General purpose upscaling' },
{ model: '4x-AnimeSharp.safetensors', description: 'Anime/illustration upscaling' },
{ model: 'NMKD-Siax_200k.safetensors', description: 'Photo-realistic upscaling' },
{ model: 'RealESRGAN_x4plus.safetensors', description: 'High quality general upscaling' }
];
for (const test of tests) {
console.log(`Testing ${test.model}...`);
console.log(`Purpose: ${test.description}`);
// Create workflow
const workflow = getUpscalingWorkflow(testImage, test.model, 1.0);
// Queue the prompt
const promptId = await queuePrompt(workflow);
console.log(`Prompt ID: ${promptId}`);
// Wait for completion
let completed = false;
let attempts = 0;
const maxAttempts = 60; // 60 seconds timeout
while (!completed && attempts < maxAttempts) {
await new Promise(resolve => setTimeout(resolve, 1000));
const history = await getHistory(promptId);
if (history && history[promptId]) {
const outputs = history[promptId].outputs;
if (outputs && outputs['5'] && outputs['5'].images && outputs['5'].images.length > 0) {
const imageName = outputs['5'].images[0].filename;
console.log(`✅ Upscaled image saved: ${imageName}`);
// Check upscaled image dimensions
const upscaledPath = `output/${imageName}`;
if (fs.existsSync(upscaledPath)) {
const upscaledDims = sizeOf(upscaledPath);
console.log(`Upscaled dimensions: ${upscaledDims.width}x${upscaledDims.height}`);
console.log(`Scale factor achieved: ${(upscaledDims.width / dimensions.width).toFixed(1)}x\n`);
}
completed = true;
}
}
attempts++;
}
if (!completed) {
console.log(`⚠️ Timeout waiting for ${test.model}\n`);
}
// Small delay between tests
await new Promise(resolve => setTimeout(resolve, 2000));
}
// Test the auto model selection
console.log('Testing automatic model selection...');
const autoTests = [
{ content: 'anime', expected: '4x-AnimeSharp.safetensors' },
{ content: 'photo', expected: 'RealESRGAN_x4plus.safetensors' },
{ content: 'general', expected: '4x-UltraSharp.safetensors' }
];
for (const test of autoTests) {
const selected = selectBestModel(test.content);
console.log(`Content type: "${test.content}" → Selected: ${selected}`);
console.log(`Expected: ${test.expected}, Match: ${selected === test.expected ? '✅' : '❌'}\n`);
}
console.log('✅ All upscaling tests completed!');
} catch (error) {
console.error('❌ Error during testing:', error);
} finally {
disconnect();
}
}
// Run the test
testUpscaling().catch(console.error);
