vast-sd-cli 0.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 14f381c408944e84b5dadeab890803dd4797958336eae022d78081273a05a20c
+  data.tar.gz: 7625775037c2ea2cdc5e6a6b943935c75ae11ba8178dab286d3210929c1c9f2a
+SHA512:
+  metadata.gz: 4595119068f2499dba27358f94bd0e4d27f1f41707f8395ee43923d344e92bfc338a7a0ce43bce7b114382d42bcbf87e4c2b775ebe8ff0ca124def350760e2a1
+  data.tar.gz: 8064e69102f3ab07d38779f64ce0483276f20e06234ebfbf1bd949af37a572449d1778d7e7b3a3374d4a72cd908c225e7b8d58df219f38bcc0834b68de419f00

data/bin/vast-sd

@@ -0,0 +1,49 @@
+#!/usr/bin/env ruby
+
+require "bundler/setup"
+
+require "dry/cli"
+
+require_relative "../lib/vast-sd/start"
+require_relative "../lib/vast-sd/config"
+
+DEFAULT_CONFIG_FILE_PATH = File.join(Dir.home, '.vast-sd.conf.json')
+
+def load_config(path)
+  if File.exists?(path)
+    JSON.parse(File.read path)
+  else
+    {}
+  end
+end
+
+module VastSd
+  module Commands
+    extend Dry::CLI::Registry
+
+    class Config < Dry::CLI::Command
+      desc "Run configuration wizard"
+
+      option :c, desc: "Configuration file path", default: DEFAULT_CONFIG_FILE_PATH
+
+      def call(c:, **)
+        VastSd::Config.new(path: c, config: load_config(c)).run
+      end
+    end
+
+    class Start < Dry::CLI::Command
+      desc "Start instance on vast.ai"
+
+      option :c, desc: "Configuration file path", default: DEFAULT_CONFIG_FILE_PATH
+
+      def call(c:, **)
+        VastSd::Start.new(config: load_config(c)).run
+      end
+    end
+
+    register "config", Config
+    register "start", Start
+  end
+end
+
+Dry::CLI.new(VastSd::Commands).call

data/lib/vast-sd/config.rb

@@ -0,0 +1,85 @@
+require 'json'
+
+require 'tty-logger'
+require 'tty-prompt'
+
+require_relative './vast'
+
+module VastSd
+  class Config
+    include Vast
+
+    def initialize(path: ,config:)
+      @path = path
+      @config = config
+
+      @logger = TTY::Logger.new
+    end
+
+    def run
+      preferred_gpus = pick_preferred_gpus
+      models = pick_models
+
+      if File.exists?(@path) && !ask_boolean("Do you want to overwrite your current configuration file? (y/n)")
+        return
+      end
+
+      File.write(@path, JSON.pretty_generate({
+        preferred_gpus: preferred_gpus,
+        models: models
+      }))
+
+      @logger.success "Configuration saved at:", @path
+    end
+
+    def pick_preferred_gpus
+      prompt = TTY::Prompt.new
+
+      offers = vast_cmd("search", "offers")
+
+      gpus_selection = offers.map do |offer|
+        offer["gpu_name"]
+      end.uniq
+
+      prompt.multi_select("Select your preferred GPU types", gpus_selection.sort, default: "RTX 4090", per_page: 20)
+    end
+
+    def pick_models
+
+      {}.tap do |models|
+        loop do
+          if !ask_boolean("Do you want to add #{models.count == 0 ? "a" : "another"} model? (y/n)")
+            break
+          end
+
+          model_kind = pick_model_kind
+
+          uri = ask_uri("Public URI of the model:").to_s
+
+          models[model_kind] ||= []
+          models[model_kind].push(uri)
+
+          @logger.success "Model added:", File.basename(uri)
+        end
+      end
+    end
+
+    def pick_model_kind
+      prompt = TTY::Prompt.new
+
+      prompt.select("What kind of model do you want to add?", ["Stable-diffusion", "Lora"])
+    end
+
+    def ask_boolean(label)
+      prompt = TTY::Prompt.new
+
+      prompt.ask(label, convert: :boolean)
+    end
+
+    def ask_uri(label)
+      prompt = TTY::Prompt.new
+
+      prompt.ask(label, convert: :uri)
+    end
+  end
+end

data/lib/vast-sd/start.rb

@@ -0,0 +1,133 @@
+require 'shellwords'
+require 'json'
+require 'tty-logger'
+require 'tty-spinner'
+require 'tty-prompt'
+
+require_relative './vast'
+
+module VastSd
+  class Start
+    include Vast
+
+    def initialize(config:)
+      @config = config
+    end
+
+    def run
+      logger = TTY::Logger.new
+      prompt = TTY::Prompt.new
+
+      user_info = vast_cmd("show", "user")
+      logger.info "Credit balance", sprintf("$%2.2f", user_info["credit"])
+
+      offers = vast_cmd("search", "offers")
+
+      offers = offers.select do |offer|
+        has_good_gpu =
+          if @config["preferred_gpus"].to_a.count > 0
+            @config["preferred_gpus"].include?(offer["gpu_name"])
+          else
+            true
+          end
+
+        has_good_gpu && offer["rentable"]
+      end.sort_by do |offer|
+        offer["dph_total"]
+      end
+
+      id = prompt.select("Instance selection", offers.map { |offer|
+        {
+          name: sprintf("%-40s %s", "#{offer['num_gpus']} #{offer["gpu_name"]} - #{offer["geolocation"]}", sprintf("$%2.2f/hr", offer["dph_total"])),
+          value: offer["id"]
+        }
+      }, per_page: 20)
+
+      contract = vast_cmd(
+        "create",
+        "instance", id,
+        "--disk", 20,
+        "--jupyter",
+        "--jupyter-dir", "/",
+        "--direct",
+        "--env", "-e JUPYTER_DIR=/ -p 3000:3000",
+        "--onstart-cmd", %{
+          sed -i '/rsync -au --remove-source-files \/venv\/ \/workspace\/venv\//a source \/workspace\/venv\/bin\/activate\n pip install jupyter_core' /start.sh;
+          /start.sh
+        },
+        "--image", "runpod/stable-diffusion:web-automatic-8.0.3"
+      )
+
+      contract_id = contract["new_contract"]
+
+      at_exit do
+        vast_cmd("destroy", "instance", contract_id)
+
+        logger.success "Instance destroyed"
+      end
+
+      logger.success "Contract ID", contract_id
+
+      # Wait until instance is ready
+      TTY::Spinner.new("[:spinner] Waiting for instance: :status", clear: true).run do |spinner|
+        loop do
+          instance = vast_cmd("show", "instance", contract_id)
+
+          spinner.update status: instance["actual_status"] || "Initialization"
+
+          if instance["actual_status"] == "running"
+            port = instance["ports"]["3000/tcp"][0]["HostPort"]
+
+            logger.info "http://#{instance["public_ipaddr"]}:#{port}"
+
+            break
+          else
+            sleep 5
+          end
+        end
+      end
+
+      ssh_info = vast_cmd('ssh-url', contract_id)
+      logger.info ssh_info
+
+      @config["models"].entries.each do |namespace, urls|
+        urls.each do |url|
+          TTY::Spinner.new("[:spinner] Loading \"#{File.basename(url)}\" in models/#{namespace}", clear: true).run do
+            upload_cmd = "curl #{Shellwords.escape(url)} -o /workspace/stable-diffusion-webui/models/#{namespace}/#{File.basename(url)}"
+
+            loop do
+              system("ssh -o StrictHostKeyChecking=no -q #{Shellwords.escape(ssh_info.chomp)} #{Shellwords.escape(upload_cmd)} &> /dev/null")
+              break if $? == 0
+            end
+          end
+        end
+      end
+
+      logger.success "Instance ready"
+
+      # Run watchdog
+      spinner = TTY::Spinner.new("[:spinner] Running: :balance")
+      spinner.auto_spin
+
+      loop do
+        begin
+          user_info = vast_cmd("show", "user")
+          spinner.update balance: "credit balance #{sprintf("$%2.2f", user_info["credit"])}"
+
+          instance = vast_cmd("show", "instance", contract_id)
+
+          if instance["actual_status"] != "running"
+            logger.error "Instance #{instance['actual_status']}"
+            exit
+          else
+            sleep 60
+          end
+        rescue Interrupt
+          exit(0)
+        end
+      end
+
+      spinner.stop
+    end
+  end
+end

data/lib/vast-sd/vast.rb

@@ -0,0 +1,33 @@
+require 'shellwords'
+require 'json'
+
+require 'tty-spinner'
+require 'tty-command'
+
+module VastSd
+  module Vast
+    def vast_cmd(*params)
+      spinner = TTY::Spinner.new(clear: true)
+      spinner.auto_spin
+    
+      cmd = TTY::Command.new(printer: :null)
+    
+      params = params.map do |param|
+        Shellwords.escape(param)
+      end
+    
+      output, error = cmd.run("vast #{params.join(" ")} --raw 2> /dev/null")
+    
+      # Cleanup output
+      output = output.gsub(/.+HHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHh/m, '')
+    
+      spinner.stop
+    
+      begin
+        JSON.parse(output)
+      rescue JSON::ParserError
+        output
+      end
+    end
+  end
+end

data/lib/vast-sd/version.rb

@@ -0,0 +1,3 @@
+module VastSd
+  VERSION = '0.1'
+end

metadata

@@ -0,0 +1,119 @@
+--- !ruby/object:Gem::Specification
+name: vast-sd-cli
+version: !ruby/object:Gem::Version
+  version: '0.1'
+platform: ruby
+authors:
+- Anonymous
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2023-08-30 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: dry-cli
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 1.0.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 1.0.0
+- !ruby/object:Gem::Dependency
+  name: tty-logger
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.6.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.6.0
+- !ruby/object:Gem::Dependency
+  name: tty-spinner
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.9.3
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.9.3
+- !ruby/object:Gem::Dependency
+  name: tty-prompt
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.23.1
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.23.1
+- !ruby/object:Gem::Dependency
+  name: tty-command
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.10.1
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 0.10.1
+description:
+email:
+- anonymous@unknown.com
+executables:
+- vast-sd
+extensions: []
+extra_rdoc_files: []
+files:
+- bin/vast-sd
+- lib/vast-sd/config.rb
+- lib/vast-sd/start.rb
+- lib/vast-sd/vast.rb
+- lib/vast-sd/version.rb
+homepage:
+licenses:
+- MIT
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 2.6.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.3.3
+signing_key:
+specification_version: 4
+summary: Little CLI tool to run and provision Stable Diffusion on Vast.ai (WIP)
+test_files: []