replicate-ruby 0.1.6 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1c51346f40acaa0f9e2f0c1ac97630fa2908d5e87d94de2d27196f2b33ec6135
-  data.tar.gz: 7559f35badbfd904058c70f3561d45793b6f3290e2905500902b755d2c45aeb2
+  metadata.gz: 82385c582239598bb4b62a3e133228b7e6fc417ea57a047c1489e5d94d787faf
+  data.tar.gz: 4ccecfdb33f1a4cf3218587c7ed4fa28475da1ca13369e59d5875b07e4f092f8
 SHA512:
-  metadata.gz: fb6c205ce5e7cd48decf8ea79067fe4eaa6942ee00c50b85a8efb82d549ae2deec969591b8ddee302b424d583d5e20f3d2c9f860e10bde8d51d633b6c9564036
-  data.tar.gz: 6d4255676a72b959ddb9ba3280adf181e6e6bbdd9012ac0085e0173bedb63e97a3458f2e8201dd22854fda4f10967ecb075f428759b67b31fd9550ad689e0fba
+  metadata.gz: c8edef78b1309f5873997b33b1de5de82f2f8982ce2ec3ed32b0cdeca0e9038c77b4d56e6173441d106160e77fb1e0d2311024d5cae2c68a1fc91ef6c62cb725
+  data.tar.gz: 56958524475fbd6636e1eb3f409b4dad100fe8179a982a746a0ee2ea4356ffa66724269a9da0e226e374bbe34c2526aacf3a358e5847ed55429ebf06700f2383

data/Gemfile.lock

@@ -1,9 +1,10 @@
 PATH
   remote: .
   specs:
-    replicate-ruby (0.1.6)
+    replicate-ruby (0.1.7)
       addressable
       faraday (>= 2.0)
+      faraday-multipart
       faraday-retry
 
 GEM
@@ -12,12 +13,15 @@ GEM
     addressable (2.8.1)
       public_suffix (>= 2.0.2, < 6.0)
     ast (2.4.2)
+    byebug (11.1.3)
     coderay (1.1.3)
     crack (0.4.5)
       rexml
     faraday (2.7.2)
       faraday-net_http (>= 2.0, < 3.1)
       ruby2_keywords (>= 0.0.4)
+    faraday-multipart (1.0.4)
+      multipart-post (~> 2)
     faraday-net_http (3.0.2)
     faraday-retry (2.0.0)
       faraday (~> 2.0)
@@ -25,6 +29,7 @@ GEM
     json (2.6.2)
     method_source (1.0.0)
     minitest (5.16.3)
+    multipart-post (2.2.3)
     parallel (1.22.1)
     parser (3.1.2.1)
       ast (~> 2.4.1)
@@ -57,9 +62,11 @@ GEM
       hashdiff (>= 0.4.0, < 2.0.0)
 
 PLATFORMS
+  arm64-darwin-22
   x86_64-darwin-21
 
 DEPENDENCIES
+  byebug
   minitest (~> 5.0)
   pry
   rake (~> 13.0)

data/README.md

@@ -54,6 +54,42 @@ id = prediction.id # store prediction id in your backend
 prediction = Replicate.client.retrieve_prediction(id) # retrieve prediction during webhook with id from backend
 ```
 
+## Dreambooth
+
+There is support for the [experimental dreambooth endpoint](https://replicate.com/blog/dreambooth-api).
+
+First, upload your training dataset:
+
+```
+upload = Replicate.client.create_upload
+upload.attach('tmp/data.zip') # replace with the path to your zip file
+```
+
+Then start training a new model using, for instance:
+
+```
+training = Replicate.client.create_training(
+  input: {
+    instance_prompt: "zwx style",
+    class_prompt: "style",
+    instance_data: upload.serving_url,
+    max_train_steps: 5000
+  },
+  model: 'yourusername/yourmodel'
+)
+```
+
+As soon as the model has finished training, you can run predictions on it:
+
+```
+prediction = Replicate.client.create_prediction(
+  input: {
+    prompt: 'your prompt, zwx style'
+  },
+  version: training.version
+)
+```
+
 ## Development
 
 After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake test` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.

data/lib/replicate/client/model.rb

@@ -9,20 +9,22 @@ module Replicate
       def retrieve_model(model, version: :latest)
         case version
         when :latest
-          Replicate::Record::Model.new(self, get("models/#{model}"))
+          response = api_endpoint.get("models/#{model}")
+          Replicate::Record::Model.new(self, response)
         when :all
-          response = get("models/#{model}/versions")
+          response = api_endpoint.get("models/#{model}/versions")
           response["results"].map! { |result| Replicate::Record::ModelVersion.new(self, result) }
           response
         else
-          Replicate::Record::ModelVersion.new(self, get("models/#{model}/versions/#{version}"))
+          response = api_endpoint.get("models/#{model}/versions/#{version}")
+          Replicate::Record::ModelVersion.new(self, response)
         end
       end
 
       # Get a collection of models
       # @see https://replicate.com/docs/reference/http#get-collection
       def retrieve_collection(slug)
-        get("collections/#{slug}")
+        api_endpoint.get("collections/#{slug}")
       end
     end
   end

data/lib/replicate/client/prediction.rb

@@ -7,13 +7,14 @@ module Replicate
       # Get a prediction
       # @see https://replicate.com/docs/reference/http#get-prediction
       def retrieve_prediction(id)
-        Replicate::Record::Prediction.new(self, get("predictions/#{id}"))
+        response = api_endpoint.get("predictions/#{id}")
+        Replicate::Record::Prediction.new(self, response)
       end
 
       # Get a list of predictions
       # @see https://replicate.com/docs/reference/http#get-predictions
       def list_predictions(cursor = nil)
-        response = get("predictions", cursor: cursor)
+        response = api_endpoint.get("predictions", cursor: cursor)
         response["results"].map! { |result| Replicate::Record::Prediction.new(self, result) }
         response
       end
@@ -22,13 +23,15 @@ module Replicate
       # @see https://replicate.com/docs/reference/http#create-prediction
       def create_prediction(params)
         params[:webhook_completed] ||= webhook_url
-        Replicate::Record::Prediction.new(self, post("predictions", params))
+        response = api_endpoint.post("predictions", params)
+        Replicate::Record::Prediction.new(self, response)
       end
 
       # Cancel a prediction
       # @see https://replicate.com/docs/reference/http#cancel-prediction
       def cancel_prediction(id)
-        Replicate::Record::Prediction.new(self, post("predictions/#{id}/cancel"))
+        response = api_endpont.post("predictions/#{id}/cancel")
+        Replicate::Record::Prediction.new(self, response)
       end
     end
   end

data/lib/replicate/client/training.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Replicate
+  class Client
+    # Methods for the Prediction API
+    module Training
+      # Get a training
+      # @see https://replicate.com/blog/dreambooth-api
+      def retrieve_training(id)
+        response = dreambooth_endpoint.get("trainings/#{id}")
+        Replicate::Record::Training.new(self, response)
+      end
+
+      # Create a training
+      # @see https://replicate.com/blog/dreambooth-api
+      def create_training(params)
+        params[:webhook_completed] ||= webhook_url
+        response = dreambooth_endpoint.post("trainings", params)
+        Replicate::Record::Training.new(self, response)
+      end
+    end
+  end
+end

data/lib/replicate/client/upload.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Replicate
+  class Client
+    # Methods for the Prediction API
+    module Upload
+      # Create an upload
+      # @see https://replicate.com/blog/dreambooth-api
+      def create_upload
+        response = dreambooth_endpoint.post("upload/data.zip")
+        Replicate::Record::Upload.new(self, response)
+      end
+
+      # Create an upload
+      # @see https://replicate.com/blog/dreambooth-api
+      def update_upload(upload_endpoint_url, zip_path)
+        endpoint = Replicate::Endpoint.new(endpoint_url: upload_endpoint_url, api_token: nil)
+        endpoint.agent.put do |req|
+          req.headers["Content-Type"] = "application/zip"
+          req.headers["Content-Length"] = File.size(zip_path).to_s
+          req.headers["Transfer-Encoding"] = "chunked"
+          req.body = Faraday::UploadIO.new(zip_path, 'application/zip')
+        end
+      end
+    end
+  end
+end

data/lib/replicate/client.rb

@@ -1,18 +1,21 @@
 # frozen_string_literal: true
 
 require "replicate/configurable"
-require "replicate/connection"
+require "replicate/endpoint"
 
 require "replicate/client/model"
 require "replicate/client/prediction"
+require "replicate/client/upload"
+require "replicate/client/training"
 
 module Replicate
   class Client
     include Replicate::Configurable
-    include Replicate::Connection
 
     include Replicate::Client::Model
     include Replicate::Client::Prediction
+    include Replicate::Client::Upload
+    include Replicate::Client::Training
 
     def initialize(options = {})
       # Use options passed in, but fall back to module defaults
@@ -21,5 +24,13 @@ module Replicate
         instance_variable_set(:"@#{key}", value)
       end
     end
+
+    def api_endpoint
+      @api_endpoint ||= Replicate::Endpoint.new(endpoint_url: api_endpoint_url, api_token: api_token)
+    end
+
+    def dreambooth_endpoint
+      @dreambooth_endpoint ||= Replicate::Endpoint.new(endpoint_url: dreambooth_endpoint_url, api_token: api_token)
+    end
   end
 end

data/lib/replicate/configurable.rb

@@ -3,13 +3,13 @@
 module Replicate
   module Configurable
     attr_accessor :api_token, :webhook_url
-    attr_writer :api_endpoint
+    attr_writer :api_endpoint_url, :dreambooth_endpoint_url
 
     class << self
       # List of configurable keys for {Datatrans::Client}
       # @return [Array] of option keys
       def keys
-        @keys ||= %i[api_token api_endpoint webhook_url]
+        @keys ||= %i[api_token api_endpoint_url dreambooth_endpoint_url webhook_url]
       end
     end
 
@@ -19,8 +19,12 @@ module Replicate
     end
 
     # API endpoint methods
-    def api_endpoint
-      @api_endpoint ||= "https://api.replicate.com/v1"
+    def api_endpoint_url
+      @api_endpoint_url ||= "https://api.replicate.com/v1"
+    end
+
+    def dreambooth_endpoint_url
+      @dreambooth_endpoint_url ||= "https://dreambooth-api-experimental.replicate.com/v1"
     end
 
     private

data/lib/replicate/{connection.rb → endpoint.rb}

@@ -3,16 +3,19 @@
 require "faraday"
 require "faraday/net_http"
 require "faraday/retry"
+require "faraday/multipart"
 require "addressable/uri"
 
 module Replicate
   # Network layer for API clients.
-  module Connection
-    DEFAULT_MEDIA_TYPE = "application/json"
-    USER_AGENT = "Datatrans Ruby Gem"
+  class Endpoint
+    attr_reader :endpoint_url, :api_token, :content_type
 
-    # Header keys that can be passed in options hash to {#get},{#head}
-    CONVENIENCE_HEADERS = Set.new(%i[accept content_type])
+    def initialize(endpoint_url:, api_token:, content_type: 'application/json')
+      @endpoint_url = endpoint_url
+      @api_token = api_token
+      @content_type = content_type
+    end
 
     # Make a HTTP GET request
     #
@@ -72,11 +75,10 @@ module Replicate
     #
     # @return [Sawyer::Agent]
     def agent
-      @agent ||= Faraday.new(url: endpoint) do |conn|
+      @agent ||= Faraday.new(url: endpoint_url) do |conn|
         conn.request :retry
-        conn.request :authorization, 'Token', api_token
-        conn.headers["Content-Type"] = DEFAULT_MEDIA_TYPE
-        conn.headers["Accept"] = DEFAULT_MEDIA_TYPE
+        conn.request :authorization, 'Token', api_token if api_token
+        conn.headers["Content-Type"] = content_type
 
         conn.adapter :net_http
       end
@@ -89,12 +91,6 @@ module Replicate
       @last_response if defined? @last_response
     end
 
-    protected
-
-    def endpoint
-      api_endpoint
-    end
-
     private
 
     def request(method, path, data)
@@ -103,7 +99,12 @@ module Replicate
       when 400
         raise Error, "#{@last_response.status} #{@last_response.reason_phrase}: #{JSON.parse(@last_response.body)}"
       else
-        JSON.parse(@last_response.body)
+        case content_type
+        when 'application/json'
+          JSON.parse(@last_response.body)
+        else
+          @last_response.body
+        end
       end
     end
   end

data/lib/replicate/record/prediction.rb

@@ -19,6 +19,26 @@ module Replicate
           false
         end
       end
+
+      def starting?
+        status == "starting"
+      end
+
+      def processing?
+        status == "processing"
+      end
+
+      def succeeded?
+        status == "succeeded"
+      end
+
+      def failed?
+        status == "failed"
+      end
+
+      def canceled?
+        status == "canceled"
+      end
     end
   end
 end

data/lib/replicate/record/training.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Replicate
+  module Record
+    class Training < Base
+      def refetch
+        @data = client.retrieve_training(id).data
+      end
+    end
+  end
+end

data/lib/replicate/record/upload.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Replicate
+  module Record
+    class Upload < Base
+      def attach(path)
+        client.update_upload(upload_url, path)
+      end
+    end
+  end
+end

data/lib/replicate/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Replicate
-  VERSION = "0.1.6"
+  VERSION = "0.2.0"
 end

data/lib/replicate.rb

@@ -7,6 +7,8 @@ require "replicate/record/base"
 require "replicate/record/model"
 require "replicate/record/model_version"
 require "replicate/record/prediction"
+require "replicate/record/upload"
+require "replicate/record/training"
 
 module Replicate
   class Error < StandardError; end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: replicate-ruby
 version: !ruby/object:Gem::Version
-  version: 0.1.6
+  version: 0.2.0
 platform: ruby
 authors:
 - Dreaming Tulpa
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2022-12-19 00:00:00.000000000 Z
+date: 2023-03-08 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: faraday
@@ -38,6 +38,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: faraday-multipart
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 - !ruby/object:Gem::Dependency
   name: addressable
   requirement: !ruby/object:Gem::Requirement
@@ -80,6 +94,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: byebug
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 description:
 email:
 - hey@dreamingtulpa.com
@@ -93,16 +121,21 @@ files:
 - Gemfile.lock
 - README.md
 - Rakefile
+- data.zip
 - lib/replicate.rb
 - lib/replicate/client.rb
 - lib/replicate/client/model.rb
 - lib/replicate/client/prediction.rb
+- lib/replicate/client/training.rb
+- lib/replicate/client/upload.rb
 - lib/replicate/configurable.rb
-- lib/replicate/connection.rb
+- lib/replicate/endpoint.rb
 - lib/replicate/record/base.rb
 - lib/replicate/record/model.rb
 - lib/replicate/record/model_version.rb
 - lib/replicate/record/prediction.rb
+- lib/replicate/record/training.rb
+- lib/replicate/record/upload.rb
 - lib/replicate/version.rb
 - sig/replicate/ruby.rbs
 homepage: https://github.com/dreamingtulpa/replicate-ruby
@@ -126,7 +159,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.26
+rubygems_version: 3.4.6
 signing_key:
 specification_version: 4
 summary: Ruby client for Replicate