link/image with pyerlai (tbc)

1 files changed, 80 insertions, 0 deletions

diff --git a/lib/plugins/link/image.ex b/lib/plugins/link/image.ex
new file mode 100644
index 0000000..cf3d9b0
--- /dev/null
+++ b/lib/plugins/link/image.ex
@@ -0,0 +1,80 @@
+defmodule Nola.Plugins.Link.Image do
+  require Logger
+  @behaviour Nola.Plugins.Link
+
+  @impl true
+  def match(_, _), do: false
+
+  @impl true
+  def post_match(_url, "image/"<>_, _header, _opts) do
+    {:body, nil}
+  end
+
+  def post_match(_, _, _, _), do: false
+
+  @impl true
+  def post_expand(_url, bytes, _, opts) do
+    pil_process = Keyword.get(opts, :pil_process, {:pil, :"py@127.0.0.1"})
+    clip_ask_process = Keyword.get(opts, :clip_ask_process, {:clip_ask, :"py@127.0.0.1"})
+    img2txt_process = Keyword.get(opts, :img2txt_process, {:image_to_text_vit_gpt2, :"py@127.0.0.1"})
+
+    tasks = [
+      Task.async(fn ->
+        {:ok, pil} = GenServer.call(pil_process, {:run, bytes})
+        pil = pil
+        |> Enum.map(fn({k, v}) -> {String.to_atom(to_string(k)), v} end)
+        pil
+      end),
+      Task.async(fn ->
+        {:ok, descr} = GenServer.call(img2txt_process, {:run, bytes})
+        {:img2txt, to_string(descr)}
+      end),
+      Task.async(fn ->
+        {:ok, prompts} = GenServer.call(clip_ask_process, {:run, bytes})
+
+        prompts = prompts
+        |> Enum.sort_by(& elem(&1, 1), &>=/2)
+        |> Enum.take(3)
+        |> Enum.map(& to_string(elem(&1, 0)))
+        |> Enum.join(", ")
+        {:prompts, prompts}
+      end)
+    ]
+    |> Task.yield_many(5000)
+    |> Enum.map(fn {task, res} ->
+      res || Task.shutdown(task, :brutal_kill)
+    end)
+
+    results = Enum.into(List.flatten(for({:ok, value} <- tasks, do: value)), Map.new)
+    img2txt = Map.get(results, :img2txt)
+    prompts = Map.get(results, :prompts)
+
+    pil = if Map.get(results, :width) do
+      animated = if Map.get(results, :animated), do: " animated", else: ""
+      "#{Map.get(results, :width, 0)}x#{Map.get(results, :height, 0)}#{animated} — "
+    else
+      ""
+    end
+
+    descr = cond do
+      img2txt && prompts ->
+        "#{pil}#{prompts} — #{img2txt}"
+      img2txt ->
+        "#{pil}#{img2txt}"
+      prompts ->
+        "#{pil}#{prompts}"
+      pil != "" ->
+        "#{pil}"
+      true ->
+        nil
+    end
+
+    if descr do
+      {:ok, "image: #{descr}"}
+    else
+      :error
+    end
+
+  end
+
+end