Class: VastSd::Config

Inherits:
Object
  • Object
show all
Includes:
Vast
Defined in:
lib/vast-sd/config.rb

Instance Method Summary collapse

Methods included from Vast

#vast_cmd

Constructor Details

#initialize(path:, config:) ⇒ Config

Returns a new instance of Config.



12
13
14
15
16
17
 
# File 'lib/vast-sd/config.rb', line 12

def initialize(path: ,config:)
  @path = path
  @config = config

  @logger = TTY::Logger.new
end

Instance Method Details

#ask_boolean(label) ⇒ Object 



73
74
75
76
77
 
# File 'lib/vast-sd/config.rb', line 73

def ask_boolean(label)
  prompt = TTY::Prompt.new

  prompt.ask(label, convert: :boolean)
end

#ask_uri(label) ⇒ Object 



79
80
81
82
83
 
# File 'lib/vast-sd/config.rb', line 79

def ask_uri(label)
  prompt = TTY::Prompt.new

  prompt.ask(label, convert: :uri)
end

#pick_model_kind ⇒ Object 



67
68
69
70
71
 
# File 'lib/vast-sd/config.rb', line 67

def pick_model_kind
  prompt = TTY::Prompt.new

  prompt.select("What kind of model do you want to add?", ["Stable-diffusion", "Lora"])
end

#pick_models ⇒ Object 



47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
 
# File 'lib/vast-sd/config.rb', line 47

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

#pick_preferred_gpus ⇒ Object 



35
36
37
38
39
40
41
42
43
44
45
 
# File 'lib/vast-sd/config.rb', line 35

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

#run ⇒ Object 



19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
 
# File 'lib/vast-sd/config.rb', line 19

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