Page MenuHomePhabricator
Files F59440

No OneTemporary
Actions

View Options

diff --git a/lib/nola/plugins.ex b/lib/nola/plugins.ex
index 7872cd6..ac94736 100644
--- a/lib/nola/plugins.ex
+++ b/lib/nola/plugins.ex
@@ -1,136 +1,137 @@
defmodule Nola.Plugins do
require Logger
@builtins [
Nola.Plugins.Account,
Nola.Plugins.Alcoolog,
Nola.Plugins.AlcoologAnnouncer,
Nola.Plugins.Base,
Nola.Plugins.Boursorama,
Nola.Plugins.Buffer,
Nola.Plugins.Calc,
Nola.Plugins.Coronavirus,
Nola.Plugins.Correction,
Nola.Plugins.Dice,
Nola.Plugins.Finance,
Nola.Plugins.Gpt,
+ Nola.Plugins.Image,
Nola.Plugins.KickRoulette,
Nola.Plugins.LastFm,
Nola.Plugins.Link,
Nola.PLugins.Logger,
Nola.Plugins.Preums,
Nola.Plugins.QuatreCentVingt,
Nola.Plugins.RadioFrance,
Nola.Plugins.Say,
Nola.Plugins.Script,
Nola.Plugins.Seen,
Nola.Plugins.Sms,
Nola.Plugins.Tell,
Nola.Plugins.Txt,
Nola.Plugins.Untappd,
Nola.Plugins.UserMention,
Nola.Plugins.WolframAlpha,
Nola.Plugins.YouTube,
]
defmodule Supervisor do
use DynamicSupervisor
require Logger
def start_link() do
DynamicSupervisor.start_link(__MODULE__, [], name: __MODULE__)
end
def start_child(module, opts \\ []) do
Logger.info("Starting #{module}")
spec = %{id: {Nola.Plugins,module}, start: {Nola.Plugins, :start_link, [module, opts]}, name: module, restart: :transient}
case DynamicSupervisor.start_child(__MODULE__, spec) do
{:ok, _} = res -> res
:ignore ->
Logger.warn("Ignored #{module}")
:ignore
{:error,_} = res ->
Logger.error("Could not start #{module}: #{inspect(res, pretty: true)}")
res
end
end
@impl true
def init(_init_arg) do
DynamicSupervisor.init(
strategy: :one_for_one,
max_restarts: 10,
max_seconds: 1
)
end
end
def dets(), do: to_charlist(Nola.data_path("/plugins.dets"))
def setup() do
:dets.open_file(dets(), [])
end
def enabled() do
:dets.foldl(fn
{name, true, _}, acc -> [name | acc]
_, acc -> acc
end, [], dets())
end
def start_all() do
Logger.info("starting plugins.")
for mod <- enabled(), do: {mod, __MODULE__.Supervisor.start_child(mod)}
end
def declare(module) do
case get(module) do
:disabled -> :dets.insert(dets(), {module, true, nil})
_ -> nil
end
end
def declare_all_builtins do
for b <- @builtins, do: declare(b)
end
def start(module, opts \\ []) do
__MODULE__.Supervisor.start_child(module)
end
@doc "Enables a plugin"
def enable(name), do: switch(name, true)
@doc "Disables a plugin"
def disable(name), do: switch(name, false)
@doc "Enables or disables a plugin"
def switch(name, value) when is_boolean(value) do
last = case get(name) do
{:ok, last} -> last
_ -> nil
end
:dets.insert(dets(), {name, value, last})
end
@spec get(module()) :: {:ok, last_start :: nil | non_neg_integer()} | :disabled
def get(name) do
case :dets.lookup(dets(), name) do
[{name, enabled, last_start}] -> {:ok, enabled, last_start}
_ -> :disabled
end
end
def start_link(module, options \\ []) do
with {:disabled, {_, true, last}} <- {:disabled, get(module)},
{:throttled, false} <- {:throttled, false}
do
module.start_link()
else
{error, _} ->
Logger.info("#{__MODULE__}: #{to_string(module)} ignored start: #{to_string(error)}")
:ignore
end
end
end
diff --git a/lib/plugins/image.ex b/lib/plugins/image.ex
new file mode 100644
index 0000000..446cb49
--- /dev/null
+++ b/lib/plugins/image.ex
@@ -0,0 +1,246 @@
+defmodule Nola.Plugins.Image do
+ require Logger
+ import Nola.Plugins.TempRefHelper
+
+ def irc_doc() do
+ """
+ # Image Generation
+
+ * **`!d2 [-n 1..10] [-g 256, 512, 1024] <prompt>`** generate image(s) using OpenAI Dall-E 2
+ * **`!sd [options] <prompt>`** generate image(s) using Stable Diffusion models (see below)
+
+ ## !sd
+
+ * `-m X` (sd2) Model (sd2: Stable Diffusion v2, sd1: Stable Diffusion v1.5, any3: Anything v3, any4: Anything v4, oj: OpenJourney)
+ * `-w X, -h X` (512) width and height. (128, 256, 384, 448, 512, 576, 640, 704, 768)
+ * `-n 1..10` (1) number of images to generate
+ * `-s X` (null) Seed
+ * `-S 0..500` (50) denoising steps
+ * `-X X` (KLMS) scheduler (DDIM, K_EULER, DPMSolverMultistep, K_EULER_ANCESTRAL, PNDM, KLMS)
+ * `-g 1..20` (7.5) guidance scale
+ * `-P 0.0..1.0` (0.8) prompt strength
+ """
+ end
+
+ def start_link() do
+ GenServer.start_link(__MODULE__, [], name: __MODULE__)
+ end
+
+ defstruct [:temprefs]
+
+ def init(_) do
+ regopts = [plugin: __MODULE__]
+ {:ok, _} = Registry.register(Nola.PubSub, "trigger:d2", regopts)
+ {:ok, _} = Registry.register(Nola.PubSub, "trigger:sd", regopts)
+ {:ok, %__MODULE__{temprefs: new_temp_refs()}}
+ end
+
+ def handle_info({:irc, :trigger, "sd", msg = %Nola.Message{trigger: %Nola.Trigger{type: :bang, args: args}}}, state) do
+ {:noreply, case OptionParser.parse(args, aliases: [m: :model], strict: [model: :string]) do
+ {_, [], _} ->
+ msg.replyfun.("#{msg.sender.nick}: sd: missing prompt")
+ state
+ {opts, prompt, _} ->
+ process_sd(Keyword.get(opts, :model, "sd2"), Enum.join(prompt, " "), msg, state)
+ end}
+ end
+
+ def handle_info({:irc, :trigger, "d2", msg = %Nola.Message{trigger: %Nola.Trigger{type: :bang, args: args}}}, state) do
+ opts = OptionParser.parse(args,
+ aliases: [n: :n, g: :geometry],
+ strict: [n: :integer, geometry: :integer]
+ )
+ case opts do
+ {_opts, [], _} ->
+ msg.replyfun.("#{msg.sender.nick}: d2: missing prompt")
+ {:noreply, state}
+ {opts, prompts, _} ->
+ prompt = Enum.join(prompts, " ")
+ geom = Keyword.get(opts, :geometry, 256)
+ request = %{
+ "prompt" => prompt,
+ "n" => Keyword.get(opts, :n, 1),
+ "size" => "#{geom}x#{geom}",
+ "response_format" => "b64_json",
+ "user" => msg.account.id,
+ }
+
+ id = FlakeId.get()
+
+ state = case OpenAi.post("/v1/images/generations", request) do
+ {:ok, %{"data" => data}} ->
+ urls = for {%{"b64_json" => b64}, idx} <- Enum.with_index(data) do
+ with {:ok, body} <- Base.decode64(b64),
+ <<smol_body::binary-size(20), _::binary>> = body,
+ {:ok, magic} <- GenMagic.Pool.perform(Nola.GenMagic, {:bytes, smol_body}),
+ bucket = Application.get_env(:nola, :s3, []) |> Keyword.get(:bucket),
+ s3path = "#{msg.account.id}/iD2#{id}#{idx}.png",
+ s3req = ExAws.S3.put_object(bucket, s3path, body, acl: :public_read, content_type: magic.mime_type),
+ {:ok, _} <- ExAws.request(s3req),
+ path = NolaWeb.Router.Helpers.url(NolaWeb.Endpoint) <> "/files/#{s3path}"
+ do
+ {:ok, path}
+ end
+ end
+
+ urls = for {:ok, path} <- urls, do: path
+ msg.replyfun.("#{msg.sender.nick}: #{Enum.join(urls, " ")}")
+ state
+ {:error, atom} when is_atom(atom) ->
+ Logger.error("dalle2: #{inspect atom}")
+ msg.replyfun.("#{msg.sender.nick}: dalle2: ☠️ #{to_string(atom)}")
+ state
+ error ->
+ Logger.error("dalle2: #{inspect error}")
+ msg.replyfun.("#{msg.sender.nick}: dalle2: ☠️ ")
+ state
+ end
+ {:noreply, state}
+ end
+ end
+
+ defp process_sd(model, prompt, msg, state) do
+ {general_opts, _, _} = OptionParser.parse(msg.trigger.args,
+ aliases: [n: :number, w: :width, h: :height],
+ strict: [number: :integer, width: :integer, height: :integer]
+ )
+
+ general_opts = general_opts
+ |> Keyword.put_new(:number, 1)
+
+ case sd_model(model, prompt, general_opts, msg.trigger.args) do
+ {:ok, env} ->
+ base_url = "https://api.runpod.ai/v1/#{env.name}"
+ {headers, options} = runpod_headers(env, state)
+ result = with {:ok, json} <- Poison.encode(%{"input" => env.request}),
+ {:ok, %HTTPoison.Response{status_code: 200, body: body}} <- HTTPoison.post("#{base_url}/run", json, headers, options),
+ {:ok, %{"id" => id} = data} <- Poison.decode(body) do
+ Logger.debug("runpod: started job #{id}: #{inspect data}")
+ spawn(fn() -> runpod_result_loop("#{base_url}/status/#{id}", env, msg, state) end)
+ :ok
+ else
+ {:ok, %HTTPoison.Response{status_code: code}} -> {:error, Plug.Conn.Status.reason_atom(code)}
+ {:error, %HTTPoison.Error{reason: reason}} -> {:error, reason}
+ end
+
+ case result do
+ {:error, reason} ->
+ Logger.error("runpod: http error for #{base_url}/run: #{inspect reason}")
+ msg.replyfun.("#{msg.sender.nick}: sd: runpod failed: #{inspect reason}")
+ _ -> :ok
+ end
+ {:error, error} ->
+ msg.replyfun.("#{msg.sender.nick}: sd: #{error}")
+ end
+
+ state
+ end
+
+ defp runpod_result_loop(url, env, msg, state) do
+ Logger.debug("runpod_result_loop: new")
+ {headers, options} = runpod_headers(env, state)
+ with {:ok, %HTTPoison.Response{status_code: 200, body: body}} <- HTTPoison.get(url, headers ++ [{"content-type", "application/json"}], options),
+ {:ok, %{"status" => "COMPLETED"} = data} <- Poison.decode(body) do
+ id = FlakeId.get()
+ tasks = for {%{"image" => url, "seed" => seed}, idx} <- Enum.with_index(Map.get(data, "output", [])) do
+ Task.async(fn() ->
+with {:ok, %HTTPoison.Response{status_code: 200, body: body}} <- HTTPoison.get(url, [], options),
+bucket = Application.get_env(:nola, :s3, []) |> Keyword.get(:bucket),
+s3path = "#{msg.account.id}/iR#{env.nick}#{id}#{idx}-#{seed}.png",
+s3req = ExAws.S3.put_object(bucket, s3path, body, acl: :public_read, content_type: "image/png"),
+{:ok, _} <- ExAws.request(s3req),
+path = NolaWeb.Router.Helpers.url(NolaWeb.Endpoint) <> "/files/#{s3path}"
+do
+ {:ok, path}
+else
+ error ->
+ Logger.error("runpod_result: error while uploading #{url}: #{inspect error}")
+ {:error, error}
+end
+ end)
+ end
+ |> Task.yield_many(5000)
+ |> Enum.map(fn {task, res} ->
+ res || Task.shutdown(task, :brutal_kill)
+ end)
+
+ results = for({:ok, {:ok, url}} <- tasks, do: url)
+
+ msg.replyfun.("#{msg.sender.nick}: #{Enum.join(results, " ")}")
+ else
+ {:ok, %{"status" => "FAILED"} = data} ->
+ Logger.error("runpod_result_loop: job FAILED: #{inspect data}")
+ msg.replyfun.("#{msg.sender.nick}: sd: job failed: #{Map.get(data, "error", "error")}")
+ {:ok, %{"status" => _} = data} ->
+ Logger.debug("runpod_result_loop: not completed: #{inspect data}")
+ :timer.sleep(:timer.seconds(1))
+ runpod_result_loop(url, env, msg, state)
+ {:ok, %HTTPoison.Response{status_code: 403}} ->
+ msg.replyfun.("#{msg.sender.nick}: sd: runpod failure: unauthorized")
+ error ->
+ Logger.warning("image: sd: runpod http error: #{inspect error}")
+ :timer.sleep(:timer.seconds(2))
+ runpod_result_loop(url, env, msg, state)
+ end
+ end
+
+ defp runpod_headers(_env, _state) do
+ config = Application.get_env(:nola, :runpod, [])
+ headers = [{"user-agent", "nola.lol bot, href@random.sh"},
+ {"authorization", "Bearer " <> Keyword.get(config, :key, "unset-api-key")}]
+ options = [timeout: :timer.seconds(180), recv_timeout: :timer.seconds(180)]
+ {headers, options}
+ end
+
+ defp sd_model(name, _, general_opts, opts) when name in ~w(sd2 sd1 oj any any4) do
+ {opts, prompt, _} = OptionParser.parse(opts, [
+ aliases: [P: :strength, s: :seed, S: :steps, g: :guidance, X: :scheduler, q: :negative],
+ strict: [strength: :float, steps: :integer, guidance: :float, scheduler: :string, seed: :integer, negative: :keep]
+ ])
+ opts = general_opts ++ opts
+ prompt = Enum.join(prompt, " ")
+
+ negative = case Keyword.get_values(opts, :negative) do
+ [] -> nil
+ list -> Enum.join(list, " ")
+ end
+
+ full_name = case name do
+ "sd2" -> "stable-diffusion-v2"
+ "sd1" -> "stable-diffusion-v1"
+ "oj" -> "sd-openjourney"
+ "any" -> "sd-anything-v3"
+ "any4" -> "sd-anything-v4"
+ end
+
+ default_scheduler = case name do
+ "sd2" -> "KLMS"
+ _ -> "K-LMS"
+ end
+
+ request = %{
+ "prompt" => prompt,
+ "num_outputs" => general_opts[:number],
+ "width" => opts[:width] || 512,
+ "height" => opts[:height] || 512,
+ "prompt_strength" => opts[:strength] || 0.8,
+ "num_inference_steps" => opts[:steps] || 30,
+ "guidance_scale" => opts[:guidance] || 7.5,
+ "scheduler" => opts[:scheduler] || default_scheduler,
+ "seed" => opts[:seed] || :rand.uniform(100_000_00)
+ }
+
+ request = if negative do
+ Map.put(request, "negative_prompt", negative)
+ else
+ request
+ end
+
+ {:ok, %{name: full_name, nick: name, request: request}}
+ end
+
+ defp sd_model(name, _, _, _) do
+ {:error, "unsupported model: \"#{name}\""}
+ end
+
+end

File Metadata

Mime Type
text/x-diff
Expires
Mon, Apr 28, 12:53 AM (1 d, 6 h)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
38874
Default Alt Text
(13 KB)

Event Timeline