mirror of
https://github.com/discourse/discourse-ai.git
synced 2025-08-05 20:53:28 +00:00
6add06af8f
This allows for 2 big features: 1. Artist can ship up to 4 prompts for image generation 2. Artist can regenerate images cause it is aware of seed This allows for iteration on images maintaining visual style
74 lines
1.8 KiB
Ruby
74 lines
1.8 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module ::DiscourseAi
|
|
module Inference
|
|
class StabilityGenerator
|
|
def self.perform!(
|
|
prompt,
|
|
width: nil,
|
|
height: nil,
|
|
api_key: nil,
|
|
engine: nil,
|
|
api_url: nil,
|
|
image_count: 4,
|
|
seed: nil
|
|
)
|
|
api_key ||= SiteSetting.ai_stability_api_key
|
|
engine ||= SiteSetting.ai_stability_engine
|
|
api_url ||= SiteSetting.ai_stability_api_url
|
|
|
|
headers = {
|
|
"Content-Type" => "application/json",
|
|
"Accept" => "application/json",
|
|
"Authorization" => "Bearer #{api_key}",
|
|
}
|
|
|
|
sdxl_allowed_dimentions = [
|
|
[1024, 1024],
|
|
[1152, 896],
|
|
[1216, 832],
|
|
[1344, 768],
|
|
[1536, 640],
|
|
[640, 1536],
|
|
[768, 1344],
|
|
[832, 1216],
|
|
[896, 1152],
|
|
]
|
|
|
|
if (!width && !height)
|
|
if engine.include? "xl"
|
|
width, height = sdxl_allowed_dimentions[0]
|
|
else
|
|
width, height = [512, 512]
|
|
end
|
|
end
|
|
|
|
payload = {
|
|
text_prompts: [{ text: prompt }],
|
|
cfg_scale: 7,
|
|
clip_guidance_preset: "FAST_BLUE",
|
|
height: width,
|
|
width: height,
|
|
samples: image_count,
|
|
steps: 30,
|
|
}
|
|
|
|
payload[:seed] = seed if seed
|
|
|
|
endpoint = "v1/generation/#{engine}/text-to-image"
|
|
|
|
response = Faraday.post("#{api_url}/#{endpoint}", payload.to_json, headers)
|
|
|
|
if response.status != 200
|
|
Rails.logger.error(
|
|
"AI stability generator failed with status #{response.status}: #{response.body}}",
|
|
)
|
|
raise Net::HTTPBadResponse
|
|
end
|
|
|
|
JSON.parse(response.body, symbolize_names: true)
|
|
end
|
|
end
|
|
end
|
|
end
|