saving prompt and revised_prompt with the PNG file

lamg · Jun 30, 2024 · b5feec4 · b5feec4
1 parent f913529
commit b5feec4
Show file tree

Hide file tree

Showing 8 changed files with 48 additions and 23 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.0.7] - 2024-06-30
+
+Added:
+
+- Images saved to the directory where `r0b0t` is running now contain the following metadata:
+  - `prompt`: the prompt used to generate the image
+  - `revised_prompt`: the revised prompt returned by `dall-e-3` after generating the image
+
 ## [0.0.6] - 2024-06-27
 
 Fixed:

diff --git a/Lib/GUI/InputOutput.fs b/Lib/GUI/InputOutput.fs
@@ -3,36 +3,48 @@ module GUI.InputOutput
 open System
 open Gdk
 open Gtk
-
 open GetProviderImpl
+open SixLabors.ImageSharp
+open SixLabors.ImageSharp.Formats.Png.Chunks
+open Stream.Types
 
 type InputOutput =
   { getPrompt: unit -> Prompt
     keyRelease: (KeyReleaseEventArgs -> unit) -> unit
     insertWord: string -> unit
-    insertImage: string -> unit }
+    insertImage: PngData -> unit }
+
+[<Literal>]
+let promptPngMetadataKey = "prompt"
+
+[<Literal>]
+let revisedPromptPngMetadataKey = "revised_prompt"
 
-let saveImage (bs: byte array) =
+let saveImage (d: PngData) =
+  use image = Image.Load d.image
+  let pngMeta = image.Metadata.GetPngMetadata()
+  pngMeta.TextData.Add(PngTextData(promptPngMetadataKey, d.prompt, "en", promptPngMetadataKey))
+  pngMeta.TextData.Add(PngTextData(revisedPromptPngMetadataKey, d.revisedPrompt, "en", revisedPromptPngMetadataKey))
   let now = DateTimeOffset.UtcNow.ToString("yyyy-MM-dd'T'HH-mm-ss")
-  IO.File.WriteAllBytes($"img_{now}.png", bs)
+  let outputPath = $"img_{now}.png"
+  image.Save(outputPath)
 
-let insertImage (insertWord: string -> unit) (chatDisplay: TextView) (content: string) =
-  let bs = Convert.FromBase64String content
-  saveImage bs
-  let p = new PixbufLoader(bs)
+let insertImage (insertWord: string -> unit) (chatDisplay: TextView) (d: PngData) =
+
+  saveImage d
+  let p = new PixbufLoader(d.image)
   let height = chatDisplay.AllocatedHeight / 2
   let width = chatDisplay.AllocatedWidth / 2
   let np = p.Pixbuf.ScaleSimple(height, width, InterpType.Bilinear)
   let buff = chatDisplay.Buffer
 
   GLib.Idle.Add(fun _ ->
     buff.InsertPixbuf(ref buff.EndIter, np)
+    insertWord d.revisedPrompt
+    insertWord "\n\n"
     false)
   |> ignore
 
-  insertWord "\n\n"
-
-
 let getPrompt (chatInput: TextView, chatDisplay: TextView) =
   fun () ->
     let question = chatInput.Buffer.Text

diff --git a/Lib/GUI/Main.fs b/Lib/GUI/Main.fs
@@ -45,7 +45,7 @@ let newAppWindow () =
       let insertData =
         function
         | Word w -> io.insertWord w
-        | PngBase64 i -> io.insertImage i
+        | PngData i -> io.insertImage i
 
       let getProvider = newGetProvider (confHandler.getConf ()) io.getPrompt
 

diff --git a/Lib/Lib.fsproj b/Lib/Lib.fsproj
@@ -39,8 +39,8 @@
     <PackageReference Include="FsHttp" Version="14.5.0" />
     <PackageReference Include="GtkSharp" Version="3.24.24.95" />
     <PackageReference Include="LamgEnv" Version="0.0.2" />
-    <PackageReference Include="LangChain" Version="0.15.0" />
     <PackageReference Include="OpenAI" Version="2.0.0-beta.3" />
+    <PackageReference Include="SixLabors.ImageSharp" Version="3.1.4" />
   </ItemGroup>
 
 </Project>
diff --git a/Lib/ProviderModuleImpl/OpenAI.fs b/Lib/ProviderModuleImpl/OpenAI.fs
@@ -24,11 +24,6 @@ let toAsyncSeqString resp =
       let last = [ None ] |> AsyncSeq.ofSeq
       AsyncSeq.append xs last
 
-let imagineFake (key: Key) (description: Prompt) =
-  let bs =
-    System.IO.File.ReadAllBytes "logo_small.png" |> System.Convert.ToBase64String
-
-  [ Some(PngBase64 bs); None ] |> AsyncSeq.ofSeq
 
 let imagine (key: Key) (description: Prompt) =
   let client = new OpenAIService(OpenAiOptions(ApiKey = key))
@@ -51,8 +46,13 @@ let imagine (key: Key) (description: Prompt) =
     let imgs =
       if resp.Successful then
         resp.Results
-        |> Seq.map (fun x -> [ x.B64 |> PngBase64 |> Some; x.RevisedPrompt |> Word |> Some ])
-        |> Seq.concat
+        |> Seq.map (fun x ->
+          { image = System.Convert.FromBase64String x.B64
+            prompt = description
+            revisedPrompt = x.RevisedPrompt }
+          |> PngData
+          |> Some)
+
       else
         resp.Error.Messages |> Seq.map (Word >> Some)
 

diff --git a/Lib/Stream/Consumer.fs b/Lib/Stream/Consumer.fs
@@ -10,7 +10,7 @@ let consumeStream (si: StopInsert) (read: Stream) =
       let! r = read ()
 
       match r with
-      | Some(Some(PngBase64 bs)) when String.IsNullOrEmpty bs -> return! loop ()
+      | Some(Some(PngData { image = [||] })) -> return! loop ()
       | Some(Some(Word w)) when String.IsNullOrEmpty w -> return! loop ()
       | Some(Some w) ->
         si.insertData w

diff --git a/Lib/Stream/Types.fs b/Lib/Stream/Types.fs
@@ -2,9 +2,14 @@ module Stream.Types
 
 open FSharp.Control
 
+type PngData =
+  { image: byte array
+    prompt: string
+    revisedPrompt: string }
+
 type LlmData =
   | Word of string
-  | PngBase64 of string
+  | PngData of PngData
 
 type GetProvider = unit -> AsyncSeq<LlmData option>
 

diff --git a/Main/Main.fsproj b/Main/Main.fsproj
@@ -9,7 +9,7 @@
     <PackageOutputPath>./nupkg</PackageOutputPath>
     <GenerateDocumentationFile>true</GenerateDocumentationFile>
 
-    <Version>0.0.6</Version>
+    <Version>0.0.7</Version>
     <Authors>Luis Ángel Méndez Gort</Authors>
     <PackageLicenseExpression>Apache-2.0</PackageLicenseExpression>
     <AssemblyName>r0b0t</AssemblyName>