Merge pull request #3 from zeke/add-support-for-MediaProvenance

add support for MediaProvenance

README.md

@@ -11,6 +11,7 @@ A Node.js CLI for generating AI images with Replicate and saving them to disk.
 - Saves all the files to disk for you.
 - Includes prediction id and a slug of the prompt in the filename.
 - Lets you specify an exact prompt, or specify a subject and roll the dice with random [promptmaker](https://npm.im/promptmaker) prompts.
+- Automatically adds [MediaProvenance](https://github.com/zeke/media-provenance) EXIF metadata to the downloaded image files so you have a record of the model, input, output, etc.
 
 ## Installation
 

lib/run.js

@@ -4,6 +4,8 @@ import dotenv from 'dotenv'
 import download from 'download'
 import { slugify } from 'transliteration'
 import path from 'node:path'
+import MediaProvenance from 'media-provenance'
+
 dotenv.config()
 
 const replicate = new Replicate()
@@ -22,18 +24,18 @@ export async function run ({ model, prompt, outputDir }) {
   const input = { prompt }
 
   let output
-  let predictionId
+  let prediction
   try {
-    output = await replicate.run(model, { input }, ({ id }) => {
-      predictionId = id
+    output = await replicate.run(model, { input }, (predictionData) => {
+      prediction = predictionData
     })
   } catch (error) {
     if (error.response?.status === 422) {
       console.log(`Received 422 error from the Replicate API. Attempting to fetch version data for ${model}...`)
       const versionedModel = await getModelFullNameWithVersion(model)
       console.log(`Rerunning with versioned model: ${versionedModel}`)
-      output = await replicate.run(versionedModel, { input }, ({ id }) => {
-        predictionId = id
+      output = await replicate.run(versionedModel, { input }, (predictionData) => {
+        prediction = predictionData
       })
     } else {
       console.error('Error running model:', error)
@@ -58,8 +60,25 @@ export async function run ({ model, prompt, outputDir }) {
   for (const url of output) {
     const urlPath = new URL(url).pathname
     const fileExtension = path.extname(urlPath)
-    const filename = `${predictionId}-${slugify(input.prompt).substring(0, 100)}${fileExtension}`
+    const filename = `${prediction.id}-${slugify(input.prompt).substring(0, 100)}${fileExtension}`
 
     await download(url, outputDir, { filename })
+
+    // Add MediaProvenance metadata to the image file
+    // See https://github.com/zeke/media-provenance
+    const fullPath = path.join(outputDir, filename)
+    const provenanceData = {
+      provider: 'Replicate (https://replicate.com/)',
+      model,
+      input: prediction.input,
+      output: prediction.output,
+      meta: {
+        ...prediction,
+        input: undefined,
+        output: undefined,
+        logs: undefined
+      }
+    }
+    MediaProvenance.set(fullPath, provenanceData)
   }
 }

package.json

@@ -16,6 +16,7 @@
   "dependencies": {
     "dotenv": "^16.4.1",
     "download": "^8.0.0",
+    "media-provenance": "^1.0.1",
     "minimist": "^1.2.8",
     "promptmaker": "^1.1.0",
     "replicate": "^0.30.2",