summaryrefslogtreecommitdiff
path: root/lib/plugins/image.ex
diff options
context:
space:
mode:
Diffstat (limited to 'lib/plugins/image.ex')
-rw-r--r--lib/plugins/image.ex246
1 files changed, 0 insertions, 246 deletions
diff --git a/lib/plugins/image.ex b/lib/plugins/image.ex
deleted file mode 100644
index 446cb49..0000000
--- a/lib/plugins/image.ex
+++ /dev/null
@@ -1,246 +0,0 @@
-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
generated by cgit v1.2.3 (git 2.25.1) at 2025-07-24 03:09:45 +0000