├── .env.example
├── .gitignore
├── LICENSE
├── README.md
├── configs
    └── sdxl.json
├── index.js
├── lib
    ├── arg-parser.js
    ├── chatgpt-model.js
    ├── image-input-model.js
    ├── replicate-model.js
    └── run-instance.js
├── loader.js
├── models
    ├── animatediff.js
    ├── bark.js
    ├── blip.js
    ├── controlnet.js
    ├── esrgan.js
    ├── gpt3.js
    ├── image-prompts.js
    ├── instructblip.js
    ├── replicate.js
    ├── sadtalker.js
    ├── sdxl.js
    ├── stablelm.js
    ├── tango.js
    ├── tortoise.js
    ├── zeroscope-upscale.js
    └── zeroscope.js
├── package-lock.json
├── package.json
├── predict.js
├── tests
    └── lib
    │   ├── arg-parser.test.js
    │   ├── replicate-model.test.js
    │   └── run-instance.test.js
└── workflows
    ├── describe.js
    └── double-upscale.js


/.env.example:
--------------------------------------------------------------------------------
1 | OPENAI_API_KEY=
2 | REPLICATE_API_TOKEN=
3 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | .env
2 | node_modules
3 | outputs/
4 | inputs/
5 | .DS_Store
6 | temp/
7 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023 fofr
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Replicate predict
 2 | 
 3 | A wrapper library to make running parallel API calls to Replicate models easier.
 4 | 
 5 | - command line models and params
 6 | - batch process prompts in a text file in an `inputs` directory
 7 | - automaticall save outputs (alongside params) using curl to an `outputs` directory
 8 | 
 9 | [See supported models](https://github.com/fofr/replicate-predict/tree/main/models)
10 | 
11 | ## Requirements
12 | 
13 | You need:
14 | 
15 | - An account on Replicate
16 | - Your Replicate API token
17 | 
18 | ## Installation
19 | 
20 | ```sh
21 | git clone https://github.com/fofr/replicate-predict
22 | cd replicate-predict
23 | npm install
24 | ```
25 | 
26 | Rename `.env.example` as `.env` and add your Replicate API token.
27 | 
28 | ## Usage
29 | 
30 | ```sh
31 | node index.js <model_name> <option> [--<option_name> <option_value>]*
32 | ```
33 | 
34 | `model_name`: The name of the model you want to use (bark, kandinsky, or stablelm).
35 | `option`: Either -a, --all, or a model’s default input, such as `prompt` or `image`.
36 | 
37 | ### Examples
38 | 
39 | ```sh
40 | node index.js bark "this world"
41 | node index.js bark --prompt "this world"
42 | node index.js bark --all # inputs/bark-prompts.txt
43 | node index.js bark "this world" --history_prompt en_speaker_1
44 | 
45 | node index.js kandinsky "picture of a cat" --steps 30
46 | node index.js stablelm "What is the capital of France?"
47 | node index.js modelscope "video of a bear"
48 | node index.js esrgan https://pbs.twimg.com/profile_images/1592472040992284674/oo_p8ahp_400x400.jpg
49 | ```
50 | 
51 | ### Example LLM output
52 | 
53 | ```sh
54 | node index.js stablelm "What is the capital of France? Have there ever been other capitals?"
55 | ```
56 | 
57 | Gives the file `outputs/stablelm/[timestamp]_What_is_the_capital_of_France.txt` with the content:
58 | 
59 | ```
60 | What is the capital of France? Have there ever been other capitals?
61 | ---
62 | The city where Napoleon Bonaparte was born is called Bordeaux. However, it’s unlikely you want me to answer this question directly without knowing what kind of “capital” means – I think we might need clarification on your definition! If you can clarify what sort of capital would you like to know about, then I can try my best to give you some useful info based on historical facts or official documents from time periods when each type of location had something unique within its jurisdiction …
63 | ```
64 | 


--------------------------------------------------------------------------------
/configs/sdxl.json:
--------------------------------------------------------------------------------
1 | {
2 |   "refine": "expert_ensemble_refiner",
3 |   "high_noise_frac": 0.95,
4 |   "apply_watermark": false,
5 |   "num_inference_steps": 35,
6 |   "num_outputs": 4
7 | }
8 | 


--------------------------------------------------------------------------------
/index.js:
--------------------------------------------------------------------------------
 1 | import Replicate from 'replicate'
 2 | import { config } from 'dotenv'
 3 | import { loadModels, loadWorkflows } from './loader.js'
 4 | import { parseArgs } from './lib/arg-parser.js'
 5 | import { runInstance } from './lib/run-instance.js'
 6 | 
 7 | config()
 8 | 
 9 | const replicate = new Replicate({
10 |   auth: process.env.REPLICATE_API_TOKEN
11 | })
12 | 
13 | async function main () {
14 |   const models = await loadModels()
15 |   const workflows = await loadWorkflows()
16 | 
17 |   const model = process.argv[2]
18 |   const args = process.argv.slice(3)
19 |   const options = parseArgs(args)
20 | 
21 |   if (!model || Object.keys(options).length === 0) {
22 |     console.log('Usage: node index.js <model> <option> [--<option_name> <option_value>]*')
23 |     console.log('Options: -a, --all, <single_input_value>, -c <count>, --count <count>')
24 |     console.log('Use -c or --count with a number to run predict that many times')
25 |     process.exit(1)
26 |   }
27 | 
28 |   const ModelClass = models[model.toLowerCase()]
29 |   const WorkflowClass = workflows[model.toLowerCase()]
30 |   const InstanceClass = ModelClass || WorkflowClass || models.replicate
31 | 
32 |   let instance
33 |   if (InstanceClass === models.replicate) {
34 |     const [user, modelName] = model.toLowerCase().split('/')
35 |     instance = new InstanceClass(replicate, { user, modelName })
36 |   } else {
37 |     instance = new InstanceClass(replicate)
38 |   }
39 | 
40 |   await instance.setModel()
41 |   await instance.setOutputDirectory()
42 | 
43 |   try {
44 |     const result = await runInstance(instance, options)
45 |     console.log('Result: ', result)
46 |   } catch (error) {
47 |     console.error(error.message)
48 |     process.exit(1)
49 |   }
50 | }
51 | 
52 | main()
53 | 


--------------------------------------------------------------------------------
/lib/arg-parser.js:
--------------------------------------------------------------------------------
 1 | const convertValue = (value) => {
 2 |   if (value.toLowerCase() === 'true') return true
 3 |   if (value.toLowerCase() === 'false') return false
 4 | 
 5 |   if (!isNaN(value) && value.trim() !== '') {
 6 |     const parsedValue = parseFloat(value)
 7 |     return Number.isInteger(parsedValue) ? parseInt(value) : parsedValue
 8 |   }
 9 | 
10 |   if (value.includes(',')) return value.split(',')
11 | 
12 |   return value
13 | }
14 | 
15 | const parseArgs = (args) => {
16 |   const options = {}
17 |   let unnamedArgIndex = 0
18 |   for (let i = 0; i < args.length; i++) {
19 |     if (args[i].startsWith('--')) {
20 |       options[args[i].slice(2)] = convertValue(args[++i])
21 |     } else {
22 |       options[`unnamedArg${unnamedArgIndex++}`] = convertValue(args[i])
23 |     }
24 |   }
25 |   return options
26 | }
27 | 
28 | export { parseArgs }
29 | 


--------------------------------------------------------------------------------
/lib/chatgpt-model.js:
--------------------------------------------------------------------------------
  1 | import { config } from 'dotenv'
  2 | import { ChatGPTAPI } from 'chatgpt'
  3 | import fs from 'fs/promises'
  4 | import path from 'path'
  5 | config()
  6 | 
  7 | const api = new ChatGPTAPI({
  8 |   apiKey: process.env.OPENAI_API_KEY
  9 | })
 10 | 
 11 | class ChatGPTModel {
 12 |   constructor (_replicate, defaultInputs = {}) {
 13 |     this.api = api
 14 |     this.defaultInputs = defaultInputs
 15 |     this.promptSplit = '\n'
 16 |     this.defaultSingleInputName = 'prompt'
 17 |     this.outputDirectory = `outputs/${this.constructor.name.toLowerCase()}`
 18 |     this.inputFilePath = `inputs/${this.constructor.name.toLowerCase()}-prompts.txt`
 19 |     this.ensureDirectoryExists(this.outputDirectory)
 20 |   }
 21 | 
 22 |   async ensureDirectoryExists (directoryPath) {
 23 |     try {
 24 |       await fs.access(directoryPath)
 25 |     } catch (error) {
 26 |       if (error.code === 'ENOENT') {
 27 |         await fs.mkdir(directoryPath, { recursive: true })
 28 |       } else {
 29 |         throw error
 30 |       }
 31 |     }
 32 |   }
 33 | 
 34 |   mergeInputWithDefaults (input) {
 35 |     return {
 36 |       ...this.defaultInputs,
 37 |       ...input
 38 |     }
 39 |   }
 40 | 
 41 |   generateFileName (text) {
 42 |     const timestamp = new Date().toISOString().replace(/[-:.]/g, '')
 43 |     if (!text) {
 44 |       return timestamp
 45 |     }
 46 | 
 47 |     const textStart = text.slice(0, 30).replace(/\s+/g, '_')
 48 |     return `${timestamp}_${textStart}`
 49 |   }
 50 | 
 51 |   async predict (input) {
 52 |     const mergedInput = this.mergeInputWithDefaults(input)
 53 |     console.log('Running', this.constructor.name, mergedInput)
 54 | 
 55 |     const options = {}
 56 |     if (mergedInput.systemMessage) {
 57 |       options.systemMessage = mergedInput.systemMessage
 58 |     }
 59 | 
 60 |     const response = await api.sendMessage(input.prompt, options)
 61 |     const message = response.text
 62 | 
 63 |     if (typeof this.saveOutputs === 'function') {
 64 |       await this.saveOutputs(response, input)
 65 |     }
 66 | 
 67 |     return this.output(message)
 68 |   }
 69 | 
 70 |   async predictMany (inputs = [], poolLimit = 5) {
 71 |     return asyncPool(poolLimit, inputs, async input => {
 72 |       return await this.predict(input)
 73 |     })
 74 |   }
 75 | 
 76 |   async saveFile (text, fileName) {
 77 |     await fs.writeFile(path.join(this.outputDirectory, fileName), text)
 78 |   }
 79 | 
 80 |   async savePrompt (text, fileName) {
 81 |     await this.saveFile(text, fileName)
 82 |   }
 83 | 
 84 |   async readPromptsFromFile (fileName) {
 85 |     const content = await fs.readFile(fileName, 'utf-8')
 86 |     const prompts = content.split(this.promptSplit)
 87 |     return prompts.filter(prompt => prompt.trim().length > 0)
 88 |   }
 89 | 
 90 |   async runAll () {
 91 |     const prompts = await this.readPromptsFromFile(this.inputFilePath)
 92 |     const inputs = prompts.map(prompt => ({ prompt: prompt.trim() }))
 93 |     await this.predictMany(inputs)
 94 |   }
 95 | 
 96 |   output (prediction) {
 97 |     return prediction
 98 |   }
 99 | 
100 |   async saveOutputs (prediction, input) {
101 |     const fileNameBase = this.generateFileName(input.prompt)
102 |     const text =
103 | `${input.prompt}
104 | ---
105 | ${this.output(prediction)}`
106 |     await this.saveFile(text, `${fileNameBase}.txt`)
107 |   }
108 | }
109 | 
110 | async function asyncPool (poolLimit, array, iteratorFn) {
111 |   const ret = []
112 |   const executing = []
113 | 
114 |   for (const item of array) {
115 |     const p = iteratorFn(item)
116 |     ret.push(p)
117 | 
118 |     if (poolLimit <= array.length) {
119 |       const e = p.finally(() => {
120 |         executing.splice(executing.indexOf(e), 1)
121 |       })
122 |       executing.push(e)
123 |       if (executing.length >= poolLimit) {
124 |         await Promise.race(executing)
125 |       }
126 |     }
127 |   }
128 | 
129 |   return Promise.all(ret)
130 | }
131 | 
132 | export default ChatGPTModel
133 | 


--------------------------------------------------------------------------------
/lib/image-input-model.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from './replicate-model.js'
 2 | 
 3 | class ImageInputModel extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs) {
 5 |     super(replicate, defaultInputs)
 6 |     this.defaultSingleInputName = 'image'
 7 |   }
 8 | 
 9 |   async predict (input) {
10 |     input.originalImage = input.image
11 |     input.image = await this.loadImageAsDataURI(input.image)
12 |     return await super.predict(input)
13 |   }
14 | }
15 | 
16 | export default ImageInputModel
17 | 


--------------------------------------------------------------------------------
/lib/replicate-model.js:
--------------------------------------------------------------------------------
  1 | import fs from 'fs/promises'
  2 | import path from 'path'
  3 | import { exec } from 'child_process'
  4 | 
  5 | class ReplicateModel {
  6 |   constructor (replicate, defaultInputs = {}, outputDirectory = null) {
  7 |     this.replicate = replicate
  8 |     this.defaultInputs = defaultInputs
  9 |     this.promptSplit = '\n'
 10 |     this.defaultSingleInputName = 'prompt'
 11 |     this.outputDirectory = outputDirectory
 12 |   }
 13 | 
 14 |   async ensureDirectoryExists (directoryPath) {
 15 |     try {
 16 |       await fs.access(directoryPath)
 17 |     } catch (error) {
 18 |       if (error.code === 'ENOENT') {
 19 |         await fs.mkdir(directoryPath, { recursive: true })
 20 |       } else {
 21 |         throw error
 22 |       }
 23 |     }
 24 |   }
 25 | 
 26 |   async setModel () {
 27 |     this.model = this.model || await this.replicate.models.get(this.user, this.modelName)
 28 | 
 29 |     if (!this.model) {
 30 |       throw new Error(`Model ${this.user}/${this.modelName} not found`)
 31 |     }
 32 | 
 33 |     if (!this.version) {
 34 |       this.version = this.model.latest_version.id
 35 |     }
 36 |   }
 37 | 
 38 |   async setOutputDirectory () {
 39 |     const today = new Date().toISOString().slice(0, 10)
 40 |     this.outputDirectory ||= `outputs/${this.modelName.toLowerCase()}/${today}`
 41 |     await this.ensureDirectoryExists(this.outputDirectory)
 42 |   }
 43 | 
 44 |   identifier () {
 45 |     return `${this.user}/${this.modelName}:${this.version}`
 46 |   }
 47 | 
 48 |   async loadConfigData (configPath) {
 49 |     try {
 50 |       return JSON.parse(await fs.readFile(configPath, 'utf-8'))
 51 |     } catch (error) {
 52 |       if (error.code === 'ENOENT') {
 53 |         console.log(`Config file ${configPath} not found. Skipping...`)
 54 |         return {}
 55 |       } else {
 56 |         throw error
 57 |       }
 58 |     }
 59 |   }
 60 | 
 61 |   async mergeInputWithDefaults (input) {
 62 |     let mergedInputs = { ...this.defaultInputs }
 63 | 
 64 |     // Parse and merge each config JSON file
 65 |     if (input.configs) {
 66 |       if (!Array.isArray(input.configs)) {
 67 |         input.configs = [input.configs]
 68 |       }
 69 | 
 70 |       for (const config of input.configs) {
 71 |         const configPath = `configs/${config}.json`
 72 |         const configData = await this.loadConfigData(configPath)
 73 |         mergedInputs = {
 74 |           ...mergedInputs,
 75 |           ...configData
 76 |         }
 77 |       }
 78 | 
 79 |       delete input.configs
 80 |     }
 81 | 
 82 |     // Merge the input params
 83 |     mergedInputs = {
 84 |       ...mergedInputs,
 85 |       ...input
 86 |     }
 87 | 
 88 |     // Do not pass the originalImage input to the model
 89 |     delete mergedInputs.originalImage
 90 | 
 91 |     return mergedInputs
 92 |   }
 93 | 
 94 |   timestamp () {
 95 |     return new Date().toISOString().replace(/[-:.]/g, '')
 96 |   }
 97 | 
 98 |   generateFileName (text) {
 99 |     if (!text) {
100 |       return this.timestamp()
101 |     }
102 | 
103 |     const textStart = text.slice(0, 30).replace(/\s+/g, '_')
104 |     return `${this.timestamp()}_${textStart}`
105 |   }
106 | 
107 |   async predict (input) {
108 |     const mergedInput = await this.mergeInputWithDefaults(input)
109 |     console.log('Running', `${this.user}/${this.modelName}`, truncateDataStrings(mergedInput))
110 | 
111 |     try {
112 |       const prediction = await this.replicate.run(this.identifier(), { input: mergedInput })
113 |       if (typeof this.saveOutputs === 'function') {
114 |         await this.saveOutputs(prediction, truncateDataStrings(mergedInput))
115 |       }
116 | 
117 |       return this.output(prediction)
118 |     } catch (error) {
119 |       console.error(error)
120 |       return false
121 |     }
122 |   }
123 | 
124 |   async predictMany (inputs = [], poolLimit = 5) {
125 |     return asyncPool(poolLimit, inputs, async input => {
126 |       return await this.predict(input)
127 |     })
128 |   }
129 | 
130 |   async predictFromFile (options = {}) {
131 |     const prompts = await this.readPromptsFromFile()
132 |     const inputs = prompts.map(prompt => ({
133 |       ...options,
134 |       [this.defaultSingleInputName]: prompt.trim()
135 |     }))
136 |     await this.predictMany(inputs)
137 |   }
138 | 
139 |   async saveFile (text, fileName) {
140 |     await fs.writeFile(path.join(this.outputDirectory, fileName), text)
141 |   }
142 | 
143 |   async saveInputAndPrediction (input, prediction, fileName) {
144 |     let defaultInput = input[this.defaultSingleInputName] || false
145 |     if (defaultInput) {
146 |       defaultInput = `${defaultInput}\n\n`
147 |     }
148 | 
149 |     await this.saveFile(JSON.stringify({ input, prediction }), `${fileName}.json`)
150 |     await this.saveFile(`${defaultInput}${JSON.stringify(input)}\n\n${JSON.stringify(prediction)}`, `${fileName}.txt`)
151 |   }
152 | 
153 |   async savePrompt (text, fileName) {
154 |     await this.saveFile(text, fileName)
155 |   }
156 | 
157 |   async readPromptsFromFile () {
158 |     const fileName = `inputs/${this.modelName.toLowerCase()}.txt`
159 |     const content = await fs.readFile(fileName, 'utf-8')
160 |     const prompts = content.split(this.promptSplit)
161 |     return prompts.filter(prompt => prompt.trim().length > 0)
162 |   }
163 | 
164 |   async saveOutputs (prediction, input) {
165 |     const fileNameBase = this.generateFileName(input[this.defaultSingleInputName])
166 |     await this.saveFilesUsingCurl(prediction, fileNameBase)
167 |     await this.saveInputAndPrediction(input, prediction, fileNameBase)
168 |   }
169 | 
170 |   async saveFilesUsingCurl (urls, fileName = null, directory) {
171 |     if (typeof urls === 'string') {
172 |       urls = [urls]
173 |     }
174 | 
175 |     for (let i = 0; i < urls.length; i++) {
176 |       const url = urls[i]
177 | 
178 |       // Check if the filename has an extension
179 |       const match = fileName.match(/(\.[\w\d_-]+)$/i)
180 | 
181 |       let indexFileName
182 |       if (match) {
183 |         // If there's an extension, insert the index before it
184 |         indexFileName = fileName.replace(match[0], `-${i}${match[0]}`)
185 |       } else {
186 |         // If there isn't, just append the index
187 |         indexFileName = `${fileName}-${i}`
188 |       }
189 |       await this.saveFileUsingCurl(url, indexFileName, directory)
190 |     }
191 |   }
192 | 
193 |   async convertFileToWebp (filename) {
194 |     const filenameAsWebp = filename.replace(/\.[\w\d_-]+$/i, '.webp')
195 |     const inputPath = path.join(this.outputDirectory, filename)
196 |     const outputPath = path.join(this.outputDirectory, filenameAsWebp)
197 | 
198 |     const command = `cwebp -q 80 "${inputPath}" -o "${outputPath}"`
199 |     return new Promise((resolve, _reject) => {
200 |       exec(command, error => {
201 |         if (error) {
202 |           console.error(error)
203 |           resolve()
204 |         } else {
205 |           resolve()
206 |         }
207 |       })
208 |     })
209 |   }
210 | 
211 |   async saveFileUsingCurl (url, fileName = null, directory = this.outputDirectory) {
212 |     return new Promise((resolve, _reject) => {
213 |       if (!fileName) {
214 |         const urlObj = new URL(url)
215 |         fileName = `${Date.now()}-${path.basename(urlObj.pathname)}`
216 |       }
217 | 
218 |       // if filename has no filetype, use url to determine it
219 |       if (!path.extname(fileName)) {
220 |         const urlObj = new URL(url)
221 |         fileName = `${fileName}${path.extname(urlObj.pathname)}`
222 |       }
223 | 
224 |       const outputPath = path.join(directory, fileName)
225 |       const curlCommand = `curl -s -L -o "${outputPath}" "${url}"`
226 | 
227 |       exec(curlCommand, error => {
228 |         if (error) {
229 |           console.error(error)
230 |           resolve()
231 |         } else {
232 |           resolve()
233 |         }
234 |       })
235 |     })
236 |   }
237 | 
238 |   output (prediction) {
239 |     return prediction
240 |   }
241 | 
242 |   async loadImageAsDataURI (inputPath) {
243 |     if (isURL(inputPath)) {
244 |       return inputPath
245 |     }
246 | 
247 |     const data = await fs.readFile(inputPath)
248 |     const base64 = data.toString('base64')
249 |     const fileExtension = path.extname(inputPath)
250 | 
251 |     switch (fileExtension) {
252 |       case '.png':
253 |         return `data:image/png;base64,${base64}`
254 |       case '.jpg':
255 |       case '.jpeg':
256 |         return `data:image/jpeg;base64,${base64}`
257 |       case '.mp4':
258 |         return `data:video/mp4;base64,${base64}`
259 |       default:
260 |         throw new Error(`Unsupported file extension: ${fileExtension}`)
261 |     }
262 |   }
263 | }
264 | 
265 | async function asyncPool (poolLimit, array, iteratorFn) {
266 |   const ret = []
267 |   const executing = []
268 | 
269 |   for (const item of array) {
270 |     const p = iteratorFn(item)
271 |     ret.push(p)
272 | 
273 |     if (poolLimit <= array.length) {
274 |       const e = p.finally(() => {
275 |         executing.splice(executing.indexOf(e), 1)
276 |       })
277 |       executing.push(e)
278 |       if (executing.length >= poolLimit) {
279 |         await Promise.race(executing)
280 |       }
281 |     }
282 |   }
283 | 
284 |   return Promise.all(ret)
285 | }
286 | 
287 | function truncateDataStrings (obj) {
288 |   const newObj = { ...obj }
289 |   for (const key in newObj) {
290 |     if (typeof newObj[key] === 'string' && newObj[key].startsWith('data:')) {
291 |       newObj[key] = newObj[key].substring(0, 100) + '...'
292 |     }
293 |   }
294 | 
295 |   return newObj
296 | }
297 | 
298 | /* eslint-disable no-new */
299 | function isURL (str) {
300 |   try {
301 |     new URL(str)
302 |     return true
303 |   } catch (_) {
304 |     return false
305 |   }
306 | }
307 | 
308 | export default ReplicateModel
309 | 


--------------------------------------------------------------------------------
/lib/run-instance.js:
--------------------------------------------------------------------------------
 1 | const cleanup = (options) => {
 2 |   const cleanedOptions = { ...options }
 3 |   delete cleanedOptions.unnamedArg0
 4 |   delete cleanedOptions.count
 5 |   return cleanedOptions
 6 | }
 7 | 
 8 | const isRunAll = (options) => {
 9 |   return options.unnamedArg0 &&
10 |     (options.unnamedArg0.toLowerCase() === '-a' ||
11 |     options.unnamedArg0.toLowerCase() === '--all')
12 | }
13 | 
14 | const runInstance = async (instance, options) => {
15 |   if (isRunAll(options)) {
16 |     return await instance.predictFromFile(cleanup({ ...options }))
17 |   } else {
18 |     const count = options.count || 1
19 |     const inputKey = instance.defaultSingleInputName
20 |     const input = options[inputKey] || options.unnamedArg0
21 | 
22 |     if (!input) {
23 |       throw new Error(`Missing input for ${inputKey}.`)
24 |     }
25 | 
26 |     const predictOptions = cleanup({
27 |       [inputKey]: input,
28 |       ...options
29 |     })
30 | 
31 |     if (count === 1) {
32 |       return await instance.predict(predictOptions)
33 |     }
34 | 
35 |     const inputs = new Array(count).fill(predictOptions)
36 |     return await instance.predictMany(inputs)
37 |   }
38 | }
39 | 
40 | export { runInstance }
41 | 


--------------------------------------------------------------------------------
/loader.js:
--------------------------------------------------------------------------------
 1 | import fs from 'fs'
 2 | import path from 'path'
 3 | import { fileURLToPath } from 'url'
 4 | 
 5 | const __dirname = path.dirname(fileURLToPath(import.meta.url))
 6 | 
 7 | async function loadModels () {
 8 |   const modelsDirectory = path.join(__dirname, 'models')
 9 |   const modelFiles = fs.readdirSync(modelsDirectory).filter(file => file.endsWith('.js'))
10 | 
11 |   const models = {}
12 |   for (const file of modelFiles) {
13 |     const modelName = path.basename(file, '.js')
14 |     const modelClass = await import(path.join(modelsDirectory, file))
15 |     models[modelName] = modelClass.default
16 |   }
17 | 
18 |   return models
19 | }
20 | 
21 | async function loadWorkflows () {
22 |   const workflowsDirectory = path.join(__dirname, 'workflows')
23 |   const workflowFiles = fs.readdirSync(workflowsDirectory).filter(file => file.endsWith('.js'))
24 | 
25 |   const workflows = {}
26 |   for (const file of workflowFiles) {
27 |     const workflowName = path.basename(file, '.js')
28 |     const workflowClass = await import(path.join(workflowsDirectory, file))
29 |     workflows[workflowName] = workflowClass.default
30 |   }
31 | 
32 |   return workflows
33 | }
34 | 
35 | export { loadModels, loadWorkflows }
36 | 


--------------------------------------------------------------------------------
/models/animatediff.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class AnimateDiff extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     seed: 0,
 6 |     steps: 25,
 7 |     guidance_scale: 7.5,
 8 |     motion_module: 'mm_sd_v15',
 9 |     negative_prompt: 'ugly, broken, distorted, garish',
10 |     path: 'realisticVisionV40_v20Novae.safetensors'
11 |   }) {
12 |     super(replicate, defaultInputs)
13 |     this.user = 'lucataco'
14 |     this.modelName = 'animate-diff'
15 |   }
16 | }
17 | 
18 | export default AnimateDiff
19 | 


--------------------------------------------------------------------------------
/models/bark.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class Bark extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs) {
 5 |     super(replicate, defaultInputs)
 6 |     this.user = 'suno-ai'
 7 |     this.modelName = 'bark'
 8 |     this.promptSplit = '\n---\n'
 9 |   }
10 | 
11 |   output (prediction) {
12 |     return prediction.audio_out
13 |   }
14 | 
15 |   async saveOutputs (prediction, input) {
16 |     const fileNameBase = this.generateFileName(input.prompt)
17 |     await this.saveFileUsingCurl(prediction.audio_out, fileNameBase)
18 |     await this.savePrompt(input.prompt, fileNameBase)
19 |   }
20 | }
21 | 
22 | export default Bark
23 | 


--------------------------------------------------------------------------------
/models/blip.js:
--------------------------------------------------------------------------------
 1 | import ImageInputModel from '../lib/image-input-model.js'
 2 | 
 3 | class Blip extends ImageInputModel {
 4 |   constructor (replicate, defaultInputs) {
 5 |     super(replicate, defaultInputs)
 6 |     this.user = 'salesforce'
 7 |     this.modelName = 'blip'
 8 |   }
 9 | 
10 |   async saveOutputs (prediction, input) {
11 |     const caption = prediction.split('Caption: ')[1]
12 |     const text =
13 | `${input.originalImage}
14 | ---
15 | ${prediction}`
16 |     await this.saveFile(text, `${this.generateFileName(caption)}.txt`)
17 |   }
18 | }
19 | 
20 | export default Blip
21 | 


--------------------------------------------------------------------------------
/models/controlnet.js:
--------------------------------------------------------------------------------
 1 | import ImageInputModel from '../lib/image-input-model.js'
 2 | 
 3 | class ControlNet extends ImageInputModel {
 4 |   constructor (replicate, defaultInputs = { structure: 'canny' }) {
 5 |     super(replicate, defaultInputs)
 6 |     this.user = 'rossjillian'
 7 |     this.modelName = 'controlnet_1-1'
 8 |   }
 9 | }
10 | 
11 | export default ControlNet
12 | 


--------------------------------------------------------------------------------
/models/esrgan.js:
--------------------------------------------------------------------------------
 1 | import ImageInputModel from '../lib/image-input-model.js'
 2 | 
 3 | class RealESRGAN extends ImageInputModel {
 4 |   constructor (replicate, defaultInputs) {
 5 |     super(replicate, defaultInputs)
 6 |     this.user = 'nightmareai'
 7 |     this.modelName = 'real-esrgan'
 8 |   }
 9 | }
10 | 
11 | export default RealESRGAN
12 | 


--------------------------------------------------------------------------------
/models/gpt3.js:
--------------------------------------------------------------------------------
1 | import ChatGPTModel from '../lib/chatgpt-model.js'
2 | export default ChatGPTModel
3 | 


--------------------------------------------------------------------------------
/models/image-prompts.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class ImagePrompts extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     max_length: 100,
 6 |     temperature: 1,
 7 |     top_p: 1,
 8 |     repetition_penalty: 1.5
 9 |   }) {
10 |     super(replicate, defaultInputs)
11 |     this.user = 'fofr'
12 |     this.modelName = 'image-prompts'
13 |   }
14 | 
15 |   async predict (input) {
16 |     return await super.predict({ prompt: `Image: ${input.prompt}` })
17 |   }
18 | 
19 |   output (prediction) {
20 |     let imagePrompt = prediction.join('').replace(/--ar.*$/m, '')
21 |     imagePrompt = imagePrompt.includes('Prompt: ') ? imagePrompt.split('Prompt: ')[1] : imagePrompt
22 |     return imagePrompt
23 |   }
24 | 
25 |   async saveOutputs (prediction, input) {
26 |     const fileNameBase = this.generateFileName(input.prompt)
27 |     const text =
28 | `${input.prompt}
29 | ---
30 | ${this.output(prediction)}`
31 |     await this.saveFile(text, `${fileNameBase}.txt`)
32 |   }
33 | }
34 | 
35 | export default ImagePrompts
36 | 


--------------------------------------------------------------------------------
/models/instructblip.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class InstructBlip extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     prompt: 'Describe this image'
 6 |   }) {
 7 |     super(replicate, defaultInputs)
 8 |     this.user = 'joehoover'
 9 |     this.modelName = 'instructblip-vicuna13b'
10 |     this.defaultSingleInputName = 'img'
11 |   }
12 | 
13 |   async predict (input) {
14 |     input.originalImage = input.img
15 |     input.img = await this.loadImageAsDataURI(input.img)
16 |     return await super.predict(input)
17 |   }
18 | 
19 |   output (prediction) {
20 |     return prediction.join('').trim()
21 |   }
22 | 
23 |   async saveOutputs (prediction, input) {
24 |     const fileNameBase = this.generateFileName(input.prompt)
25 |     const text =
26 | `${input.prompt}
27 | ---
28 | ${this.output(prediction)}`
29 |     await this.saveFile(text, `${fileNameBase}.txt`)
30 |   }
31 | }
32 | 
33 | export default InstructBlip
34 | 


--------------------------------------------------------------------------------
/models/replicate.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class Replicate extends ReplicateModel {
 4 |   constructor (replicate, { defaultInputs, user, modelName }) {
 5 |     super(replicate, defaultInputs)
 6 |     this.user = user
 7 |     this.modelName = modelName
 8 |   }
 9 | }
10 | 
11 | export default Replicate
12 | 


--------------------------------------------------------------------------------
/models/sadtalker.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class SadTalker extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs) {
 5 |     super(replicate, defaultInputs)
 6 |     this.user = 'cjwbw'
 7 |     this.modelName = 'sadtalker'
 8 |     this.defaultSingleInputName = 'source_image'
 9 |   }
10 | 
11 |   async saveOutputs (prediction, input) {
12 |     const fileNameBase = this.generateFileName(input.prompt)
13 |     await this.saveFileUsingCurl(prediction, fileNameBase)
14 |   }
15 | }
16 | 
17 | export default SadTalker
18 | 


--------------------------------------------------------------------------------
/models/sdxl.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class SDXL extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     num_outputs: 1,
 6 |     width: 1024,
 7 |     height: 1024,
 8 |     scheduler: 'K_EULER',
 9 |     refine: 'expert_ensemble_refiner',
10 |     high_noise_frac: 0.8,
11 |     negative_prompt: 'soft, blurry, ugly, broken, distorted, garish'
12 |   }) {
13 |     super(replicate, defaultInputs)
14 |     this.user = 'stability-ai'
15 |     this.modelName = 'sdxl'
16 |   }
17 | }
18 | 
19 | export default SDXL
20 | 


--------------------------------------------------------------------------------
/models/stablelm.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class StableLM extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs) {
 5 |     super(replicate, defaultInputs)
 6 |     this.user = 'stability-ai'
 7 |     this.modelName = 'stablelm-tuned-alpha-7b'
 8 |   }
 9 | 
10 |   output (prediction) {
11 |     return prediction.join('')
12 |   }
13 | 
14 |   async saveOutputs (prediction, input) {
15 |     const fileNameBase = this.generateFileName(input.prompt)
16 |     const text =
17 | `${input.prompt}
18 | ---
19 | ${this.output(prediction)}`
20 |     await this.saveFile(text, `${fileNameBase}.txt`)
21 |   }
22 | }
23 | 
24 | export default StableLM
25 | 


--------------------------------------------------------------------------------
/models/tango.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class Tango extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     steps: 100,
 6 |     guidance: 3
 7 |   }) {
 8 |     super(replicate, defaultInputs)
 9 |     this.user = 'cjwbw'
10 |     this.modelName = 'tango'
11 |   }
12 | }
13 | 
14 | export default Tango
15 | 


--------------------------------------------------------------------------------
/models/tortoise.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class Tortoise extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     voice_a: 'random',
 6 |     voice_b: 'random',
 7 |     voice_c: 'random'
 8 |   }) {
 9 |     super(replicate, defaultInputs)
10 |     this.user = 'afiaka87'
11 |     this.modelName = 'tortoise-tts'
12 |     this.promptSplit = '\n---\n'
13 |     this.defaultSingleInputName = 'text'
14 |   }
15 | }
16 | 
17 | export default Tortoise
18 | 


--------------------------------------------------------------------------------
/models/zeroscope-upscale.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class ZeroScopeUpscale extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     num_frames: 24,
 6 |     width: 1024,
 7 |     height: 576,
 8 |     fps: 8,
 9 |     guidance_scale: 12.5,
10 |     num_inference_steps: 50,
11 |     model: 'xl',
12 |     init_weight: 0.25,
13 |     prompt: 'detailed, 8k, beautiful',
14 |     negative_prompt: 'noisy, washed out, ugly, distorted, low quality, garish'
15 |   }) {
16 |     super(replicate, defaultInputs)
17 |     this.user = 'anotherjesse'
18 |     this.modelName = 'zeroscope-v2-xl'
19 |     this.defaultSingleInputName = 'init_video'
20 |   }
21 | 
22 |   async predict (input) {
23 |     if (input.init_video) {
24 |       input.init_video = await this.loadImageAsDataURI(input.init_video)
25 |     }
26 |     return await super.predict(input)
27 |   }
28 | }
29 | 
30 | export default ZeroScopeUpscale
31 | 


--------------------------------------------------------------------------------
/models/zeroscope.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | 
 3 | class ZeroScope extends ReplicateModel {
 4 |   constructor (replicate, defaultInputs = {
 5 |     num_frames: 24,
 6 |     fps: 8,
 7 |     width: 576,
 8 |     height: 320,
 9 |     guidance_scale: 12.5,
10 |     num_inference_steps: 50,
11 |     model: '576w',
12 |     negative_prompt: 'noisy, washed out, ugly, distorted, low quality, garish'
13 |   }) {
14 |     super(replicate, defaultInputs)
15 |     this.user = 'anotherjesse'
16 |     this.modelName = 'zeroscope-v2-xl'
17 |   }
18 | }
19 | 
20 | export default ZeroScope
21 | 


--------------------------------------------------------------------------------
/package.json:
--------------------------------------------------------------------------------
 1 | {
 2 |   "name": "replicate-predict",
 3 |   "version": "1.0.0",
 4 |   "description": "Run predictions for favourite models",
 5 |   "main": "index.js",
 6 |   "type": "module",
 7 |   "scripts": {
 8 |     "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js && npm run lint",
 9 |     "lint": "standard",
10 |     "lint:fix": "standard --fix",
11 |     "predict": "node ./index.js"
12 |   },
13 |   "author": "",
14 |   "license": "MIT",
15 |   "dependencies": {
16 |     "chatgpt": "^5.2.4",
17 |     "dotenv": "^16.0.3",
18 |     "replicate": "^0.11.1"
19 |   },
20 |   "devDependencies": {
21 |     "eslint-plugin-jest": "^27.2.1",
22 |     "jest": "^29.5.0",
23 |     "standard": "^17.0.0"
24 |   },
25 |   "jest": {
26 |     "testMatch": [
27 |       "**/tests/**/*.test.js"
28 |     ],
29 |     "transform": {}
30 |   }
31 | }
32 | 


--------------------------------------------------------------------------------
/predict.js:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env node
2 | import './index.js'
3 | 


--------------------------------------------------------------------------------
/tests/lib/arg-parser.test.js:
--------------------------------------------------------------------------------
 1 | /* eslint-env jest */
 2 | import { parseArgs } from '../../lib/arg-parser.js'
 3 | 
 4 | describe('parseArgs', () => {
 5 |   test('No arguments', () => {
 6 |     const args = []
 7 |     const expectedResult = {}
 8 | 
 9 |     expect(parseArgs(args)).toEqual(expectedResult)
10 |   })
11 | 
12 |   test('One unnamed argument', () => {
13 |     const args = ['input1']
14 |     const expectedResult = {
15 |       unnamedArg0: 'input1'
16 |     }
17 | 
18 |     expect(parseArgs(args)).toEqual(expectedResult)
19 |   })
20 | 
21 |   test('Multiple unnamed arguments', () => {
22 |     const args = ['input1', 'input2']
23 |     const expectedResult = {
24 |       unnamedArg0: 'input1',
25 |       unnamedArg1: 'input2'
26 |     }
27 | 
28 |     expect(parseArgs(args)).toEqual(expectedResult)
29 |   })
30 | 
31 |   test('Named arguments', () => {
32 |     const args = ['--option1', 'value1', '--option2', 'value2']
33 |     const expectedResult = {
34 |       option1: 'value1',
35 |       option2: 'value2'
36 |     }
37 | 
38 |     expect(parseArgs(args)).toEqual(expectedResult)
39 |   })
40 | 
41 |   test('Mixed arguments', () => {
42 |     const args = ['input1', '--option1', 'value1', 'input2', '--option2', 'value2']
43 |     const expectedResult = {
44 |       unnamedArg0: 'input1',
45 |       option1: 'value1',
46 |       unnamedArg1: 'input2',
47 |       option2: 'value2'
48 |     }
49 | 
50 |     expect(parseArgs(args)).toEqual(expectedResult)
51 |   })
52 | 
53 |   test('Arguments with numbers', () => {
54 |     const args = ['input1', '--option1', '42', '3.14', '--option2', 'value2']
55 |     const expectedResult = {
56 |       unnamedArg0: 'input1',
57 |       option1: 42,
58 |       unnamedArg1: 3.14,
59 |       option2: 'value2'
60 |     }
61 | 
62 |     expect(parseArgs(args)).toEqual(expectedResult)
63 |   })
64 | })
65 | 


--------------------------------------------------------------------------------
/tests/lib/replicate-model.test.js:
--------------------------------------------------------------------------------
  1 | /* eslint-env jest */
  2 | import ReplicateModel from '../../lib/replicate-model.js'
  3 | import fs from 'fs/promises'
  4 | import path, { dirname } from 'path'
  5 | import { jest } from '@jest/globals'
  6 | import { fileURLToPath } from 'url'
  7 | 
  8 | const __filename = fileURLToPath(import.meta.url)
  9 | const __dirname = dirname(__filename)
 10 | 
 11 | class MockReplicate {
 12 |   async run () {
 13 |     return 'Mock prediction'
 14 |   }
 15 | }
 16 | 
 17 | describe('ReplicateModel', () => {
 18 |   const replicate = new MockReplicate()
 19 |   const model = new ReplicateModel(replicate)
 20 | 
 21 |   afterEach(async () => {
 22 |     try {
 23 |       await fs.rmdir('outputs/mockmodel', { recursive: true })
 24 |     } catch (error) {
 25 |       // Ignore errors while cleaning up the output directory
 26 |     }
 27 |   })
 28 | 
 29 |   test('should create output directory', async () => {
 30 |     const outputPath = path.join(model.outputDirectory)
 31 |     await model.ensureDirectoryExists(outputPath)
 32 | 
 33 |     const stats = await fs.stat(outputPath)
 34 |     expect(stats.isDirectory()).toBe(true)
 35 |   })
 36 | 
 37 |   test('should merge input with defaults', () => {
 38 |     const input = { prompt: 'Test' }
 39 |     const mergedInput = model.mergeInputWithDefaults(input)
 40 | 
 41 |     expect(mergedInput).toEqual({ ...model.defaultInputs, ...input })
 42 |   })
 43 | 
 44 |   test('should predict', async () => {
 45 |     const input = { prompt: 'Test' }
 46 |     const prediction = await model.predict(input)
 47 | 
 48 |     expect(prediction).toBe('Mock prediction')
 49 |   })
 50 | 
 51 |   test('should read prompts from file', async () => {
 52 |     const testInputPath = 'inputs/test-inputs.txt'
 53 |     const testInputContent = 'Test1\nTest2\n\nTest3'
 54 |     await fs.writeFile(testInputPath, testInputContent)
 55 | 
 56 |     const prompts = await model.readPromptsFromFile(testInputPath)
 57 | 
 58 |     expect(prompts).toEqual(['Test1', 'Test2', 'Test3'])
 59 | 
 60 |     await fs.unlink(testInputPath)
 61 |   })
 62 | 
 63 |   test('should generate file name with or without text', () => {
 64 |     const fileNameWithText = model.generateFileName('Test input')
 65 |     expect(fileNameWithText).toMatch(/^\d{8}T\d{6}\d{3}Z_Test_input/)
 66 | 
 67 |     const fileNameWithoutText = model.generateFileName()
 68 |     expect(fileNameWithoutText).toMatch(/^\d{8}T\d{6}\d{3}Z/)
 69 |   })
 70 | 
 71 |   test('should save file', async () => {
 72 |     const filePath = path.join(model.outputDirectory, 'test-output.txt')
 73 |     await model.saveFile('Test content', 'test-output.txt')
 74 | 
 75 |     const content = await fs.readFile(filePath, 'utf-8')
 76 |     expect(content).toBe('Test content')
 77 | 
 78 |     await fs.unlink(filePath)
 79 |   })
 80 | 
 81 |   test('should save prompt', async () => {
 82 |     const filePath = path.join(model.outputDirectory, 'test-prompt.txt')
 83 |     await model.savePrompt('Test prompt', 'test-prompt.txt')
 84 | 
 85 |     const content = await fs.readFile(filePath, 'utf-8')
 86 |     expect(content).toBe('Test prompt')
 87 | 
 88 |     await fs.unlink(filePath)
 89 |   })
 90 | 
 91 |   test('should run all prompts from file', async () => {
 92 |     const testInputPath = 'inputs/test-inputs.txt'
 93 |     const testInputContent = 'Test1\nTest2\n\nTest3'
 94 |     await fs.writeFile(testInputPath, testInputContent)
 95 | 
 96 |     const originalInputFilePath = model.inputFilePath
 97 |     model.inputFilePath = testInputPath
 98 | 
 99 |     const predictManySpy = jest.spyOn(model, 'predictMany')
100 |     const predictSpy = jest.spyOn(model, 'predict')
101 |     await model.runAll()
102 | 
103 |     expect(predictManySpy).toHaveBeenCalled()
104 |     expect(predictSpy).toHaveBeenCalledTimes(3)
105 | 
106 |     model.inputFilePath = originalInputFilePath
107 |     await fs.unlink(testInputPath)
108 |   })
109 | 
110 |   test('should save file using curl', async () => {
111 |     const testImageUrl = 'https://via.placeholder.com/150'
112 |     const testOutputFileName = 'test-image.jpg'
113 |     const outputPath = path.join(model.outputDirectory, testOutputFileName)
114 | 
115 |     await model.saveFileUsingCurl(testImageUrl, testOutputFileName)
116 | 
117 |     const stats = await fs.stat(outputPath)
118 |     expect(stats.isFile()).toBe(true)
119 | 
120 |     await fs.unlink(outputPath)
121 |   })
122 | 
123 |   test('should return output as is', () => {
124 |     const testOutput = { result: 'Test output' }
125 |     const output = model.output(testOutput)
126 | 
127 |     expect(output).toEqual(testOutput)
128 |   })
129 | 
130 |   test('should load image as Data URI', async () => {
131 |     const imagePath = path.join(__dirname, 'test-image.jpg')
132 |     const imageData = Buffer.from('iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAACXBIWXMAAAsTAAALEwEAmpwYAAAAB3RJTUUH4wcOEBsE_g0_WgAAAARnQU1BAACxjwv8YQUAAAAJcEhZcwAADsMAAA7DAcdvqGQAAAAMSURBVBhXY3jy5AkABmEDcAqA8b9XAAAAAElFTkSuQmCC', 'base64')
133 |     await fs.writeFile(imagePath, imageData)
134 | 
135 |     const imageDataURI = await model.loadImageAsDataURI(imagePath)
136 | 
137 |     expect(imageDataURI).toMatch(/^data:image\/jpeg;base64,/)
138 | 
139 |     // Test with a URL
140 |     const imageUrl = 'https://via.placeholder.com/150'
141 |     const urlDataURI = await model.loadImageAsDataURI(imageUrl)
142 | 
143 |     expect(urlDataURI).toBe(imageUrl)
144 | 
145 |     await fs.unlink(imagePath)
146 |   })
147 | 
148 |   test('should predict many with varying poolLimit', async () => {
149 |     const inputs = [
150 |       { prompt: 'Test1' },
151 |       { prompt: 'Test2' },
152 |       { prompt: 'Test3' }
153 |     ]
154 | 
155 |     const predictions = await model.predictMany(inputs, 2)
156 |     expect(predictions).toHaveLength(inputs.length)
157 |     expect(predictions).toEqual(['Mock prediction', 'Mock prediction', 'Mock prediction'])
158 | 
159 |     const predictionsNoLimit = await model.predictMany(inputs)
160 |     expect(predictionsNoLimit).toHaveLength(inputs.length)
161 |     expect(predictionsNoLimit).toEqual(['Mock prediction', 'Mock prediction', 'Mock prediction'])
162 |   })
163 | })
164 | 


--------------------------------------------------------------------------------
/tests/lib/run-instance.test.js:
--------------------------------------------------------------------------------
 1 | /* eslint-env jest */
 2 | import { runInstance } from '../../lib/run-instance.js'
 3 | 
 4 | class MockInstance {
 5 |   constructor (shouldThrow) {
 6 |     this.shouldThrow = shouldThrow
 7 |   }
 8 | 
 9 |   async runAll () {
10 |     if (this.shouldThrow) {
11 |       throw new Error('Error in runAll')
12 |     }
13 |     return 'All run'
14 |   }
15 | 
16 |   async predict (options) {
17 |     if (this.shouldThrow) {
18 |       console.error('Error in predict')
19 |       return false
20 |     }
21 |     return options
22 |   }
23 | 
24 |   async predictMany (inputs = [], poolLimit = 5) {
25 |     return Promise.all(inputs.map(async input => {
26 |       return await this.predict(input)
27 |     }))
28 |   }
29 | 
30 |   get defaultSingleInputName () {
31 |     return 'input'
32 |   }
33 | }
34 | 
35 | describe('runInstance', () => {
36 |   test('should run all with --all', async () => {
37 |     const instance = new MockInstance(false)
38 |     const options = { unnamedArg0: '--all' }
39 |     const result = await runInstance(instance, options)
40 | 
41 |     expect(result).toBe('All run')
42 |   })
43 | 
44 |   test('should run predict', async () => {
45 |     const instance = new MockInstance(false)
46 |     const options = { input: 'test' }
47 |     const result = await runInstance(instance, options)
48 | 
49 |     expect(result).toEqual({ input: 'test' })
50 |   })
51 | 
52 |   test('should throw error when input is missing', async () => {
53 |     const instance = new MockInstance(false)
54 |     const options = {}
55 | 
56 |     await expect(runInstance(instance, options)).rejects.toThrow('Missing input for input')
57 |   })
58 | 
59 |   test('should handle errors in runAll', async () => {
60 |     const instance = new MockInstance(true)
61 |     const options = { unnamedArg0: '--all' }
62 | 
63 |     await expect(runInstance(instance, options)).rejects.toThrow('Error in runAll')
64 |   })
65 | 
66 |   test('should handle errors in predict', async () => {
67 |     const instance = new MockInstance(true)
68 |     const options = { input: 'test' }
69 | 
70 |     const result = await runInstance(instance, options)
71 |     expect(result).toEqual(false)
72 |   })
73 | 
74 |   test('should run predict 3 times with --count', async () => {
75 |     const instance = new MockInstance(false)
76 |     const options = { '--count': 3, input: 'test' }
77 |     const result = await runInstance(instance, options)
78 | 
79 |     expect(result).toEqual([{ input: 'test' }, { input: 'test' }, { input: 'test' }])
80 |   })
81 | 
82 |   test('should run predict 2 times with -c', async () => {
83 |     const instance = new MockInstance(false)
84 |     const options = { '-c': 2, input: 'test' }
85 |     const result = await runInstance(instance, options)
86 | 
87 |     expect(result).toEqual([{ input: 'test' }, { input: 'test' }])
88 |   })
89 | 
90 |   test('should handle errors in predict with -c', async () => {
91 |     const instance = new MockInstance(true)
92 |     const options = { '-c': 2, input: 'test' }
93 |     const result = await runInstance(instance, options)
94 | 
95 |     expect(result).toEqual([false, false])
96 |   })
97 | })
98 | 


--------------------------------------------------------------------------------
/workflows/describe.js:
--------------------------------------------------------------------------------
 1 | import ReplicateModel from '../lib/replicate-model.js'
 2 | import InstructBlip from '../models/instructblip.js'
 3 | import GPT3 from '../models/gpt3.js'
 4 | 
 5 | class Describe extends ReplicateModel {
 6 |   constructor (replicate, defaultInputs = {}) {
 7 |     super(replicate, defaultInputs)
 8 |     this.defaultSingleInputName = 'img'
 9 |     this.instructBlip = new InstructBlip(this.replicate, {
10 |       prompt: 'Give a detailed description of the image.',
11 |       max_length: 1024,
12 |       length_penalty: 5
13 |     })
14 |     this.gpt3 = new GPT3(replicate, {
15 |       systemMessage: `
16 | You will be given up to 4 descriptions of images. For each you need to convert it into the format:
17 | 
18 | [subject], [in the style of], [...keywords]
19 | 
20 | For example:
21 | - an indian mask in the jungle, in the style of seapunk, urban culture exploration, sfumato, primitive figurines, vacation dadcore, melds mexican and american cultures, tupinipunk
22 | - the animated astronaut is staring upward in space, in the style of colorful caricature, zbrush, skottie young, jeff koons, play with light, gorecore, ferrania p30
23 | - ok ka logo on a black background a gold letter logo icon design ok ka and wave with a circle behind it royalty free illustration, in the style of fine lines, delicate curves, orientalist scenes, leather/hide, poignant
24 | - a drawing of a beast has its feet stretched out, in the style of barroco, controlled chaos, sharp/prickly, outlandish energy, cryptid academia, m42 mount, unilalianism
25 | 
26 | Match the format shown above. Return them as numbered list items corresponding with the input order.
27 | `
28 |     })
29 |   }
30 | 
31 |   async predict (input) {
32 |     const blipOutputs = await Promise.all([
33 |       this.instructBlip.predict(input),
34 |       this.instructBlip.predict(input),
35 |       this.instructBlip.predict(input),
36 |       this.instructBlip.predict(input)
37 |     ])
38 | 
39 |     const joinedOutputNumbered = blipOutputs.map((output, index) => {
40 |       return `${index + 1}. ${output}`
41 |     }).join('\n\n')
42 | 
43 |     return await this.gpt3.predict({ prompt: joinedOutputNumbered })
44 |   }
45 | }
46 | 
47 | export default Describe
48 | 


--------------------------------------------------------------------------------
/workflows/double-upscale.js:
--------------------------------------------------------------------------------
 1 | import ZeroScopeUpscale from '../models/zeroscope-upscale.js'
 2 | import ReplicateModel from '../lib/replicate-model.js'
 3 | 
 4 | class DoubleUpscale extends ReplicateModel {
 5 |   constructor (replicate, defaultInputs = {}) {
 6 |     super(replicate, defaultInputs)
 7 |     this.defaultSingleInputName = 'url'
 8 |     this.zeroscope = new ZeroScopeUpscale(this.replicate)
 9 |   }
10 | 
11 |   async predict (input) {
12 |     const zOutput = await this.zeroscope.predict({ init_video: input.url })
13 |     if (zOutput && zOutput[0]) {
14 |       await this.zeroscope.predict({ init_video: zOutput[0] })
15 |     }
16 |   }
17 | }
18 | 
19 | export default DoubleUpscale
20 | 


--------------------------------------------------------------------------------