Image Style Translator

async function processImage(originalImg, styleName = 'udnie') {
    /**
     * Applies a pre-trained neural style transfer model to an image.
     * This function uses Google's Magenta.js library to restyle the content of
     * the input image using the artistic style of another image, represented by a model.
     * 
     * @param {Image} originalImg - The input Image object to be stylized.
     * @param {string} [styleName='udnie'] - The name of the style model to use.
     *   Available styles: 'udnie', 'wave', 'scream', 'rain_princess', 'la_muse', 'the_shipwreck_of_the_minotaur'.
     * @returns {Promise<HTMLCanvasElement>} A canvas element with the stylized image, or an error canvas on failure.
     */

    // Helper function to dynamically load a script if it's not already present.
    const loadScript = (url) => {
        return new Promise((resolve, reject) => {
            // Check if the script is already on the page
            if (document.querySelector(`script[src="${url}"]`)) {
                return resolve();
            }
            const script = document.createElement('script');
            script.src = url;
            script.onload = () => resolve();
            script.onerror = () => reject(new Error(`Script load error for ${url}`));
            document.head.appendChild(script);
        });
    };

    try {
        // A list of known, working style models from Magenta.js
        const availableStyles = ['udnie', 'wave', 'scream', 'rain_princess', 'la_muse', 'the_shipwreck_of_the_minotaur'];
        if (!availableStyles.includes(styleName)) {
            console.warn(`Invalid style: '${styleName}'. Defaulting to 'udnie'. Available styles: ${availableStyles.join(', ')}`);
            styleName = 'udnie';
        }

        // 1. Load required libraries dynamically.
        // Magenta.js depends on TensorFlow.js, so we load tf.js first.
        const tfjsUrl = 'https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@2.8.6/dist/tf.min.js';
        const magentaUrl = 'https://cdn.jsdelivr.net/npm/@magenta/image@1.2.1/dist/magentaimage.min.js';
        
        // Wait for TensorFlow.js to load
        if (!window.tf) {
          await loadScript(tfjsUrl);
        }
        // Wait for Magenta.js to load
        if (!window.mm) {
          await loadScript(magentaUrl);
        }

        // 2. Construct the model URL from the style name
        const modelURL = `https://storage.googleapis.com/magentadata/js/checkpoints/style/${styleName}/`;

        // 3. Initialize the style transfer model
        const styler = new mm.ImageStylization(modelURL);
        await styler.initialize();

        // 4. Create an output canvas and apply the style transfer
        // The stylize function takes an input image/canvas and returns a new canvas.
        const resultCanvas = await styler.stylize(originalImg);
        
        // 5. Clean up the model to free up GPU memory
        styler.dispose();

        // 6. Return the resulting canvas
        return resultCanvas;

    } catch (error) {
        console.error("Image style transfer failed:", error);
        
        // Return a canvas with an error message if something goes wrong
        const errorCanvas = document.createElement('canvas');
        errorCanvas.width = originalImg.naturalWidth > 0 ? originalImg.naturalWidth : 400;
        errorCanvas.height = originalImg.naturalHeight > 0 ? originalImg.naturalHeight : 200;
        const ctx = errorCanvas.getContext('2d');
        
        ctx.fillStyle = '#F0F0F0';
        ctx.fillRect(0, 0, errorCanvas.width, errorCanvas.height);
        
        ctx.fillStyle = '#D9534F'; // red
        ctx.font = 'bold 16px sans-serif';
        ctx.textAlign = 'center';
        ctx.textBaseline = 'middle';

        const lines = [
            'Style transfer failed.',
            'This can happen on the first run as the model loads.',
            'Please try again or check the console for errors.'
        ];
        const lineHeight = 22;
        const startY = errorCanvas.height / 2 - (lines.length - 1) * lineHeight / 2;

        lines.forEach((line, index) => {
            ctx.fillText(line, errorCanvas.width / 2, startY + index * lineHeight);
        });

        return errorCanvas;
    }
}