summaryrefslogtreecommitdiff
path: root/lib/plugins/link/image.ex
blob: cf3d9b0b3610b44d5ba5599306b10d15e289710a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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