From 262ed4753ee688cd62f96bcb69a12709f61dd91f Mon Sep 17 00:00:00 2001
From: Rafael dos Santos Silva <xfalcox@gmail.com>
Date: Fri, 19 May 2023 20:38:08 -0300
Subject: [PATCH] FEATURE: Basic StableDiffusion text2img support (#72)

---
 config/settings.yml                         |  7 ++++
 lib/shared/inference/stability_generator.rb | 36 +++++++++++++++++++++
 plugin.rb                                   |  1 +
 3 files changed, 44 insertions(+)
 create mode 100644 lib/shared/inference/stability_generator.rb

diff --git a/config/settings.yml b/config/settings.yml
index 0a6a21a1..ab3d90c7 100644
--- a/config/settings.yml
+++ b/config/settings.yml
@@ -94,6 +94,13 @@ plugins:
   ai_anthropic_api_key:
     default: ""
     secret: true
+  ai_stability_api_key:
+    default: ""
+    secret: true
+  ai_stability_api_url:
+    default: "https://api.stability.ai"
+  ai_stability_engine:
+    default: "stable-diffusion-xl-beta-v2-2-2"
 
   composer_ai_helper_enabled:
     default: false
diff --git a/lib/shared/inference/stability_generator.rb b/lib/shared/inference/stability_generator.rb
new file mode 100644
index 00000000..ead46f7d
--- /dev/null
+++ b/lib/shared/inference/stability_generator.rb
@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module ::DiscourseAi
+  module Inference
+    class StabilityGenerator
+      def self.perform!(prompt)
+        headers = {
+          "Referer" => Discourse.base_url,
+          "Content-Type" => "application/json",
+          "Accept" => "application/json",
+          "Authorization" => "Bearer #{SiteSetting.ai_stability_api_key}",
+        }
+
+        payload = {
+          text_prompts: [{ text: prompt }],
+          cfg_scale: 7,
+          clip_guidance_preset: "FAST_BLUE",
+          height: 512,
+          width: 512,
+          samples: 4,
+          steps: 30,
+        }
+
+        base_url = SiteSetting.ai_stability_api_url
+        engine = SiteSetting.ai_stability_engine
+        endpoint = "v1/generation/#{engine}/text-to-image"
+
+        response = Faraday.post("#{base_url}/#{endpoint}", payload.to_json, headers)
+
+        raise Net::HTTPBadResponse if response.status != 200
+
+        JSON.parse(response.body, symbolize_names: true)
+      end
+    end
+  end
+end
diff --git a/plugin.rb b/plugin.rb
index 6d6bb82f..3ab95a48 100644
--- a/plugin.rb
+++ b/plugin.rb
@@ -32,6 +32,7 @@ after_initialize do
   require_relative "lib/shared/inference/openai_completions"
   require_relative "lib/shared/inference/openai_embeddings"
   require_relative "lib/shared/inference/anthropic_completions"
+  require_relative "lib/shared/inference/stability_generator"
 
   require_relative "lib/shared/classificator"
   require_relative "lib/shared/post_classificator"