lib/gen_ai/image/stability_ai.rb in gen-ai-0.2.0 vs lib/gen_ai/image/stability_ai.rb in gen-ai-0.2.1

@@ -6,10 +6,11 @@
   class Image
     class StabilityAI < Base
       DEFAULT_SIZE = '256x256'
       API_BASE_URL = 'https://api.stability.ai'
       DEFAULT_MODEL = 'stable-diffusion-xl-beta-v2-2-2'
+      UPSCALE_MODEL = 'stable-diffusion-x4-latent-upscaler'
 
       def initialize(token:, options: {})
         build_client(token)
       end
 
@@ -37,10 +38,23 @@
           model: model,
           parsed: response['artifacts'].map { |artifact| artifact['base64'] }
         )
       end
 
+      def upscale(image, options = {})
+        model = options[:model] || UPSCALE_MODEL
+        url = "/v1/generation/#{model}/image-to-image/upscale"
+
+        response = client.post url, build_upscale_body(image, options), multipart: true
+
+        build_result(
+          raw: response,
+          model: model,
+          parsed: response['artifacts'].map { |artifact| artifact['base64'] }
+        )
+      end
+
       private
 
       def build_client(token)
         @client = GenAI::Api::Client.new(url: API_BASE_URL, token: token)
       end
@@ -54,13 +68,23 @@
           width: w
         }.merge(options)
       end
 
       def build_edit_body(image, prompt, options)
-        {
+        params = {
           init_image: File.binread(image),
           'text_prompts[0][text]' => prompt,
           'text_prompts[0][weight]' => 1.0
+        }
+        params.merge!(mask: File.binread(options.delete(:mask))) if options[:mask]
+        params.merge(options)
+      end
+
+      def build_upscale_body(image, options)
+        w, = size(options)
+        {
+          image: File.binread(image),
+          width: w
         }.merge(options)
       end
 
       def size(options)
         size = options.delete(:size) || DEFAULT_SIZE