From 827fd653fc90cb200614d8b550a017a6c5c22c67 Mon Sep 17 00:00:00 2001
From: Tijs Zwinkels <tijs@tinkertank.eu>
Date: Tue, 15 Aug 2023 15:07:37 +0200
Subject: [PATCH] Add image generattion parameters to jpeg exif metadata

---
 package.json                    |  1 +
 src/components/ImageActions.vue |  2 +-
 src/utils/base64.ts             | 27 ++++++++++++++++++++++-----
 src/utils/download.ts           | 18 +++++++++++++++---
 4 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/package.json b/package.json
index 724da13c..2a5c7727 100644
--- a/package.json
+++ b/package.json
@@ -20,6 +20,7 @@
     "jszip": "^3.10.1",
     "marked": "^4.1.1",
     "pako": "^2.1.0",
+    "piexifjs": "^1.0.6",
     "pinia": "^2.0.21",
     "rxjs": "^7.8.0",
     "sanitize-html": "^2.7.2",
diff --git a/src/components/ImageActions.vue b/src/components/ImageActions.vue
index 6286a2fe..4fb62775 100644
--- a/src/components/ImageActions.vue
+++ b/src/components/ImageActions.vue
@@ -110,7 +110,7 @@ const ratingDialog = ref(false);
 
 <template>
     <el-button @click="confirmDelete" type="danger" :icon="Delete" plain>Delete</el-button>
-    <el-button @click="downloadImage(imageData.image, `${imageData.seed}-${imageData.prompt}`)" type="success" :icon="Download" plain>Download</el-button>
+    <el-button @click="downloadImage(imageData.image, `${imageData.seed}-${imageData.prompt}`, imageData)" type="success" :icon="Download" plain>Download</el-button>
     <el-button v-if="!imageData.starred" @click="outputStore.toggleStarred(imageData.id)" type="warning" :icon="Star" plain>Favourite</el-button>
     <el-button v-if="imageData.starred" @click="outputStore.toggleStarred(imageData.id)" type="warning" :icon="StarFilled" plain>Unfavourite</el-button>
     <el-button @click="store.generateText2Img(imageData)" type="success" :icon="Refresh" plain>Text2img</el-button>
diff --git a/src/utils/base64.ts b/src/utils/base64.ts
index e94c811d..144d1b47 100644
--- a/src/utils/base64.ts
+++ b/src/utils/base64.ts
@@ -1,3 +1,4 @@
+import piexif from 'piexifjs';
 
 /**
  * Converts base64 data into to another data type
@@ -23,7 +24,7 @@ export async function convertBase64ToDataType(base64Image: string, contentType:
 /**
  * Converts base64 data into a Uint8Array
  */
-export async function convertBase64ToUint8Array(base64Image: string, contentType?: string) {
+export async function convertBase64ToUint8Array(base64Image: string, contentType?: string, userComment?: string) {
     // Split into two parts
     const parts = base64Image.split(';base64,');
 
@@ -31,7 +32,7 @@ export async function convertBase64ToUint8Array(base64Image: string, contentType
     const imageType = contentType ?? parts[0].split(':')[1];
 
     // Decode Base64 string
-    const decodedData =
+    let decodedData =
         window.atob(
             imageType === parts[0].split(':')[1] ?
                 parts[1] :
@@ -39,6 +40,22 @@ export async function convertBase64ToUint8Array(base64Image: string, contentType
         );
 
 
+    // Add exif data if filetype is jpeg, and userComment is set
+    if (imageType == "image/jpeg" && userComment) {
+        const zeroth = {
+          [piexif.ImageIFD.Software]: "Stable UI - Create images with Stable Diffusion using the AI Horde"
+        };
+        const exif = userComment ? {
+          [piexif.ExifIFD.UserComment]: `ASCII\0\0\0${userComment}`
+        } : undefined;
+        const exifObj = {
+          "0th": zeroth,
+          "Exif": exif
+        };
+        const exifbytes = piexif.dump(exifObj);
+        decodedData = piexif.insert(exifbytes, decodedData);
+    }
+
     // Create UNIT8ARRAY of size same as row data length
     const uInt8Array = new Uint8Array(decodedData.length);
 
@@ -53,12 +70,12 @@ export async function convertBase64ToUint8Array(base64Image: string, contentType
 /**
  * Converts base64 data into a blob
  */
-export async function convertBase64ToBlob(base64Image: string, contentType?: string) {
+export async function convertBase64ToBlob(base64Image: string, contentType?: string, userComment?: string) {
     // Hold the content type
     const imageType = contentType ?? base64Image.split(';base64,')[0].split(':')[1];
-    
+
     // Get Uint8Array
-    const uInt8Array = await convertBase64ToUint8Array(base64Image, contentType)
+    const uInt8Array = await convertBase64ToUint8Array(base64Image, contentType, userComment)
 
     // Return BLOB image after conversion
     return new Blob([uInt8Array], { type: imageType });
diff --git a/src/utils/download.ts b/src/utils/download.ts
index 269a33f0..4cb88119 100644
--- a/src/utils/download.ts
+++ b/src/utils/download.ts
@@ -4,6 +4,7 @@ import { convertBase64ToBlob } from './base64';
 import { downloadZip, type InputWithSizeMeta } from 'client-zip';
 import { db } from '@/utils/db';
 import type { IndexableType } from 'dexie';
+import { type ImageData } from '@/stores/outputs';
 
 export async function downloadMultipleImages(outputKeys: IndexableType[], showMessage = true, onProcess?: () => void) {
     const optionsStore = useOptionsStore();
@@ -32,7 +33,7 @@ export async function downloadMultipleImages(outputKeys: IndexableType[], showMe
         } else if (optionsStore.imageDownloadType === "JPG") {
             toDownload.push({
                 name: fileName + ".jpg",
-                input: await convertBase64ToBlob(image, "image/jpeg"),
+                input: await convertBase64ToBlob(image, "image/jpeg", getMetaData(output)),
             })
         } else {
             toDownload.push({
@@ -57,7 +58,18 @@ export async function downloadMultipleImages(outputKeys: IndexableType[], showMe
     downloadLink.click();
 }
 
-export async function downloadImage(base64Data: string, fileName: string) {
+// Create a string with the image params for the file
+function getMetaData(imageData: ImageData) : string{
+    const prompts = imageData.prompt.split("###");
+    const prompt = prompts[0].trim();
+    const negativePrompt = prompts.length > 1 ? prompts[1].trim() : null;
+    return `${prompt}\n` +
+    ((negativePrompt) ? `Negative prompt: ${negativePrompt}\n` : ``) +
+    `Steps: ${imageData.steps}, Sampler: ${imageData.sampler_name}, CFG scale: ${imageData.cfg_scale}, Seed: ${imageData.seed},` +
+    `Size: ${imageData.width}x${imageData.height}, model: ${imageData.modelName}`
+}
+
+export async function downloadImage(base64Data: string, fileName: string, imageData: ImageData) {
     const optionsStore = useOptionsStore();
 
     const downloadLink = document.createElement("a");
@@ -68,7 +80,7 @@ export async function downloadImage(base64Data: string, fileName: string) {
         downloadLink.href = URL.createObjectURL(blob);
         downloadLink.download = fileName.replace(/[/\\:*?"<>]/g, "").substring(0, 128).trimEnd() + ".png"; // Only get first 128 characters so we don't break the max file name limit
     } else if (optionsStore.imageDownloadType === "JPG") {
-        blob = await convertBase64ToBlob(base64Data, "image/jpeg");
+        blob = await convertBase64ToBlob(base64Data, "image/jpeg", getMetaData(imageData));
         downloadLink.href = URL.createObjectURL(blob);
         downloadLink.download = fileName.replace(/[/\\:*?"<>]/g, "").substring(0, 128).trimEnd() + ".jpg";
     } else {