fav

#!/usr/bin/env node

const {spawn} = require('child_process')
const glob = require('glob')
const path = require('path')
const fs = require('fs-extra')
const yargs = require('yargs') // https://www.npmjs.com/package/yargs
const VERSION = require('./package.json').version

const help = () => {
    console.log(`\nfav v${VERSION}`)
    console.log(`\nTo stylize a video:`)
    console.log(`    $ fav stylize ./path/to/video.mp4 model_file_name.t7`)
    console.log(`\nView documentation for stylize:`)
    console.log(`    $ fav stylize --help`)
    console.log(`\nTo get a list of pretrained models:`)
    console.log(`    $ fav list-models`)
    console.log('\n')
}

const startAPI = (argv) => {
    const port = argv.port || process.env.PORT
    console.log(`Starting fav API on port ${port}`)
    require('http')
        .createServer()
        .listen(port)
        .on('error', (err) => {
            console.error(err)
        })
}

yargs
    .command('$0', 'Start the API', (yargs) => {
        yargs.option('p', {
            alias: 'port',
            default: 3030,
        })
    }, startAPI)
    .command('help', 'help', () => {}, help)
    .command('list-models', 'list available pre-trained style models', () => {}, (argv) => {
        const models = glob.sync('/fast-artistic-videos/models/*video.t7')
        models.map(m => console.log(path.parse(m).base))
    })
    .command('stylize <video_file_in_cwd> <model_file_name>', 'apply style to a video', (yargs) => {
        yargs.positional('video_file_in_cwd', {
            description: 'Video file to stylize (needs to be in current working directory or lower).',
        })
        yargs.positional('model_file_name', {
            description: 'Name of the model preset to use. Get a list of options with fav list-models',
        })

        yargs.option('r', {
            alias: ['res', 'resolution'],
            description: 'Sets the resolution of the output format. e.g. 640:-1. Defaults to resolution of input.',
            default: '-1:-1',
        })
        .option('g', {
            alias: 'gpu',
            description: 'Sets the GPU index. Use -1 for CPU.',
            default: 0,
            choices: [-1, 0, 1, 2, '...'],
        })
        .option('b', {
            alias: 'backend',
            description: 'Sets the backend to use.',
            default: 'cudnn',
            choices: ['cudnn', 'nn', 'opencl']
        })
        .option('o', {
            alias: 'overwrite',
            description: 'Overwrite existing output file?',
            default: false,
            type: 'boolean'
        })
        .option('k', {
            alias: 'keep',
            description: 'Keep the process files. Keep if you want to run the same video with different style models.',
            default: false,
            type: 'boolean',
        })
    }, (argv) => {
        const spawnArgs = [
            'bash',
            ['./stylizeVideo_flownet.sh', `/io/${argv.video_file_in_cwd}`, `./models/${argv.model_file_name}`],
            {cwd: '/fast-artistic-videos'}
        ]
        // console.log(argv, spawnArgs)
        
        // check if output file already exists and handle with the --overwrite flag
        // delete the file if overwrite is true
        // kill the process if file exists and overwrite is false
        try{
            const vid = path.parse(argv.video_file_in_cwd)
            const outputPath = `/io/${vid.name}-stylized${vid.ext}`
            fs.statSync(outputPath)
            // file exists
            if(argv.overwrite === true){
                // overwrite
                fs.remove(outputPath)
            }else{
                // don't overwrite
                console.log(outputPath, 'exists. Use --overwrite to overwrite it.')
                return process.exit()
            }
        }catch(e){
            // file doesn't exist. nothing to do.
        }

        const proc = spawn(...spawnArgs)

        // accept defaults for all interactive prompts
        proc.stdin.write(argv.gpu.toString() + '\n') // which gpu
        proc.stdin.write(argv.backend + '\n') // which backend
        proc.stdin.write(argv.resolution + '\n') // resolution
        proc.stdout.on('data', (str) => console.log(str.toString()))
        proc.stderr.on('data', (str) => console.error(str.toString()))
        proc.on('close', (code) => {
            console.log(`exited with code ${code}`)
            // clean up based on flag. 
            // User would want to keep flow files if they are running the same video with multiple models
            if(!argv.keep) {
                fs.remove(`/io/${path.parse(argv.video_file_in_cwd).name}/`)
            }
        })
    })
    .argv