Class: RubyLLM::Protocols::InvokeModel::StabilityImages

Inherits:
InvokeModel
  • Object
show all
Defined in:
lib/ruby_llm/protocols/invoke_model/stability_images.rb

Overview

Stability image generation and inpainting over Bedrock InvokeModel.

Constant Summary collapse

GENERATION_MODELS = 
%w[
  stability.sd3-5-large-v1:0 stability.stable-image-core-v1:1 stability.stable-image-ultra-v1:1
].freeze
INPAINT_MODEL = 
'stability.stable-image-inpaint-v1:0'
MODELS = 
(GENERATION_MODELS + [INPAINT_MODEL]).freeze
ASPECT_RATIOS = 
%w[16:9 1:1 21:9 2:3 3:2 4:5 5:4 9:16 9:21].freeze
IMAGE_TYPES = 
%w[image/jpeg image/png image/webp].freeze

Instance Method Summary collapse

Instance Method Details

#paint(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {}) ⇒ Object 



16
17
18
19
20
21
22
23
24
 
# File 'lib/ruby_llm/protocols/invoke_model/stability_images.rb', line 16

def paint(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {})
  track_usage(:image) do
    payload = render_image_payload(prompt, model:, size:, count:, with:, mask:, provider_options:)
    response = signed_post("/model/#{model}/invoke", payload)
    images = parse_image_responses(response, model:)
    images.each { |image| image.config = @config }
    images.one? ? images.first : images
  end
end

#parse_image_responses(response, model:) ⇒ Object 



33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
 
# File 'lib/ruby_llm/protocols/invoke_model/stability_images.rb', line 33

def parse_image_responses(response, model:)
  images = Array(response.body['images']).compact.reject(&:empty?)
  if images.empty?
    reasons = Array(response.body['finish_reasons']).compact.join(', ')
    message = reasons.empty? ? 'Bedrock returned no images' : "Bedrock image generation failed: #{reasons}"
    raise Error.new(message, response:)
  end

  images.map do |data|
    mime_type = RubyLLM::Files::MimeType.for(StringIO.new(Base64.decode64(data)))
    raise Error.new('Bedrock returned invalid image data', response:) unless IMAGE_TYPES.include?(mime_type)

    Image.new(data:, model:, mime_type:)
  end
end

#render_image_payload(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {}) ⇒ Object 



26
27
28
29
30
31
 
# File 'lib/ruby_llm/protocols/invoke_model/stability_images.rb', line 26

def render_image_payload(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {})
  RubyLLM.logger.debug { 'Stability image models return one image per request' } if count && count > 1
  attachments = Attachment.wrap(with, config: @config)
  payload = { prompt: }.merge(render_image_inputs(attachments, model:, mask:, size:))
  payload.merge(provider_options.transform_keys(&:to_sym))
end