lib/gen_ai/image/open_ai.rb in gen-ai-0.2.3 vs lib/gen_ai/image/open_ai.rb in gen-ai-0.3.0

@@ -1,7 +1,9 @@
 # frozen_string_literal: true
 
+require 'base64'
+
 module GenAI
   class Image
     class OpenAI < Base
       DEFAULT_SIZE = '256x256'
       DEFAULT_MODEL = 'dall-e-3'
@@ -19,35 +21,35 @@
         response = handle_errors { @client.images.generate(parameters: parameters) }
 
         build_result(
           raw: response,
           model: parameters[:model],
-          parsed: response['data'].map { |datum| datum[RESPONSE_FORMAT] }
+          parsed: parse_response_data(response['data'])
         )
       end
 
       def variations(image, options = {})
         parameters = build_variations_options(image, options)
 
         response = handle_errors { @client.images.variations(parameters: parameters) }
 
         build_result(
           raw: response,
-          model: 'dall-e',
-          parsed: response['data'].map { |datum| datum[RESPONSE_FORMAT] }
+          model: parameters[:model],
+          parsed: parse_response_data(response['data'])
         )
       end
 
       def edit(image, prompt, options = {})
         parameters = build_edit_options(image, prompt, options)
 
         response = handle_errors { @client.images.edit(parameters: parameters) }
 
         build_result(
           raw: response,
-          model: 'dall-e',
-          parsed: response['data'].map { |datum| datum[RESPONSE_FORMAT] }
+          model: parameters[:model],
+          parsed: parse_response_data(response['data'])
         )
       end
 
       private
 
@@ -62,20 +64,26 @@
 
       def build_variations_options(image, options)
         {
           image: image,
           size: options.delete(:size) || DEFAULT_SIZE,
+          model: 'dall-e-2', # variation is only available on dall-e-2
           response_format: options.delete(:response_format) || RESPONSE_FORMAT
         }.merge(options)
       end
 
       def build_edit_options(image, prompt, options)
         {
           image: image,
           prompt: prompt,
           size: options.delete(:size) || DEFAULT_SIZE,
+          model: 'dall-e-2', # edit is only available on dall-e-2
           response_format: options.delete(:response_format) || RESPONSE_FORMAT
         }.merge(options)
+      end
+
+      def parse_response_data(data)
+        data.map { |datum| Base64.decode64(datum[RESPONSE_FORMAT]) }
       end
     end
   end
 end