discourse-ai/spec/shared/inference/stability_generator_spec.rb

50 lines
1.7 KiB
Ruby

# frozen_string_literal: true
describe DiscourseAi::Inference::StabilityGenerator do
def gen(prompt)
DiscourseAi::Inference::StabilityGenerator.perform!(prompt)
end
it "sets dimensions to 512x512 for non XL model" do
SiteSetting.ai_stability_engine = "stable-diffusion-v1-5"
SiteSetting.ai_stability_api_url = "http://www.a.b.c"
SiteSetting.ai_stability_api_key = "123"
stub_request(:post, "http://www.a.b.c/v1/generation/stable-diffusion-v1-5/text-to-image")
.with do |request|
json = JSON.parse(request.body)
expect(json["text_prompts"][0]["text"]).to eq("a cow")
expect(json["width"]).to eq(512)
expect(json["height"]).to eq(512)
expect(request.headers["Authorization"]).to eq("Bearer 123")
expect(request.headers["Content-Type"]).to eq("application/json")
true
end
.to_return(status: 200, body: "{}", headers: {})
gen("a cow")
end
it "sets dimensions to 1024x1024 for XL model" do
SiteSetting.ai_stability_engine = "stable-diffusion-xl-1024-v1-0"
SiteSetting.ai_stability_api_url = "http://www.a.b.c"
SiteSetting.ai_stability_api_key = "123"
stub_request(
:post,
"http://www.a.b.c/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image",
)
.with do |request|
json = JSON.parse(request.body)
expect(json["text_prompts"][0]["text"]).to eq("a cow")
expect(json["width"]).to eq(1024)
expect(json["height"]).to eq(1024)
expect(request.headers["Authorization"]).to eq("Bearer 123")
expect(request.headers["Content-Type"]).to eq("application/json")
true
end
.to_return(status: 200, body: "{}", headers: {})
gen("a cow")
end
end