105 lines
3.0 KiB
Elixir
105 lines
3.0 KiB
Elixir
defmodule Diffuser.Generator.PromptRequestWorker do
|
|
alias Diffuser.Generator
|
|
alias Diffuser.Generator.PromptRequest
|
|
alias DiffuserWeb.Endpoint
|
|
alias Diffuser.Repo
|
|
|
|
@path [:code.priv_dir(:diffuser), "python"] |> Path.join()
|
|
|
|
def start(%PromptRequest{} = prompt_request) do
|
|
with {:ok, active_prompt} <-
|
|
update_and_broadcast_progress(prompt_request, "in_progress"),
|
|
{:ok, _file_location} <- call_python(:test_script, :test_func, active_prompt),
|
|
%PromptRequest{} = prompt_request_with_results <-
|
|
write_and_save_images(active_prompt),
|
|
{:ok, completed_prompt} <-
|
|
update_and_broadcast_progress(prompt_request_with_results, "finished") do
|
|
{:ok, completed_prompt |> Repo.preload(:images)}
|
|
else
|
|
nil ->
|
|
raise("prompt not found")
|
|
|
|
{:error, message} ->
|
|
raise(message)
|
|
end
|
|
end
|
|
|
|
defp update_and_broadcast_progress(prompt_request, "in_progress"),
|
|
do:
|
|
update_and_broadcast_progress(prompt_request, %{
|
|
status: "in_progress",
|
|
began_at: NaiveDateTime.utc_now()
|
|
})
|
|
|
|
defp update_and_broadcast_progress(prompt_request, "finished"),
|
|
do:
|
|
update_and_broadcast_progress(prompt_request, %{
|
|
status: "finished",
|
|
ended_at: NaiveDateTime.utc_now()
|
|
})
|
|
|
|
defp update_and_broadcast_progress(%PromptRequest{id: id} = prompt_request, attrs) do
|
|
{:ok, new_prompt} = Generator.update_prompt_request(prompt_request, attrs)
|
|
|
|
:ok = Endpoint.broadcast("request:#{id}", "request", %{prompt_request: new_prompt})
|
|
|
|
{:ok, new_prompt}
|
|
end
|
|
|
|
defp call_python(_module, _func, %PromptRequest{
|
|
id: prompt_id,
|
|
prompt: prompt,
|
|
steps: steps,
|
|
guidance_scale: guidance_scale
|
|
}) do
|
|
port =
|
|
Port.open(
|
|
{:spawn,
|
|
~s(python #{@path}/stable_diffusion.py --prompt "#{prompt}" --output "#{@path}/#{prompt_id}.png" --num-inference-steps #{steps} --guidance-scale #{guidance_scale})},
|
|
[:binary]
|
|
)
|
|
|
|
python_loop(port, prompt_id)
|
|
end
|
|
|
|
defp python_loop(port, prompt_id) do
|
|
receive do
|
|
{^port, {:data, ":finished" <> msg}} ->
|
|
{:ok, msg}
|
|
|
|
{^port, {:data, ":step" <> step}} ->
|
|
{:ok, prompt_request} =
|
|
Generator.update_prompt_request(prompt_id, %{completed_steps: step})
|
|
|
|
Endpoint.broadcast(
|
|
"request:#{prompt_id}",
|
|
"progress",
|
|
%{prompt_request: prompt_request |> Repo.preload(:images)}
|
|
)
|
|
|
|
python_loop(port, prompt_id)
|
|
|
|
{^port, _result} ->
|
|
python_loop(port, prompt_id)
|
|
end
|
|
end
|
|
|
|
defp write_and_save_images(%PromptRequest{id: id}) do
|
|
file_path = "#{@path}/#{id}.png"
|
|
|
|
with {:ok, body} <- File.read(file_path),
|
|
{:ok, _result} <-
|
|
Generator.create_prompt_request_result(
|
|
id,
|
|
%{
|
|
file_name: "#{id}.png",
|
|
filename: "#{id}.png",
|
|
binary: body
|
|
}
|
|
),
|
|
:ok <- File.rm(file_path) do
|
|
Generator.get_prompt_request!(id)
|
|
end
|
|
end
|
|
end
|