styling, require a code for prompts, admin privs, ordering, filtering, jump-to-page

lib/diffuser/generator/prompt_request.ex

@@ -1,17 +1,26 @@
 defmodule Diffuser.Generator.PromptRequest do
-  use Ecto.Schema
+  use Diffuser.Schema
   use Waffle.Ecto.Schema
   import Ecto.Changeset
+  import Ecto.Query
   alias Diffuser.Generator.PromptRequestResult
+  alias Diffuser.Accounts.{User, Vote}
+
+  @valid_codes Application.get_env(:diffuser, :valid_codes, [])
 
-  @primary_key {:id, Ecto.UUID, autogenerate: true}
   schema "prompt_requests" do
     field :prompt, :string
     field :status, :string, default: "queued"
-    field :steps, :integer
-    field :guidance_scale, :float
+    field :steps, :integer, default: 32
+    field :completed_steps, :integer, default: 0
+    field :guidance_scale, :float, default: 7.5
+    field :began_at, :utc_datetime
+    field :ended_at, :utc_datetime
+    field :code, :string
 
     has_many :images, PromptRequestResult, on_delete: :delete_all
+    has_many :votes, Vote, on_delete: :delete_all
+    belongs_to :user, User
 
     timestamps()
   end
@@ -19,7 +28,63 @@ defmodule Diffuser.Generator.PromptRequest do
   @doc false
   def changeset(prompt_request, attrs) do
     prompt_request
-    |> cast(attrs, [:prompt, :status, :steps, :guidance_scale])
+    |> cast(attrs, [
+      :prompt,
+      :status,
+      :steps,
+      :guidance_scale,
+      :completed_steps,
+      :began_at,
+      :ended_at,
+      :code
+    ])
+    |> validate_number(:steps, less_than_or_equal_to: 60)
+    |> validate_code()
     |> validate_required([:prompt])
   end
+
+  defp validate_code(changeset) do
+    validate_change(changeset, :code, fn :code, code ->
+      cond do
+        @valid_codes
+        |> Enum.any?(fn {_name, valid_code} ->
+          valid_code == code
+        end) ->
+          []
+
+        true ->
+          [code: "Code is invalid."]
+      end
+    end)
+  end
+
+  def search_by_prompt_or_user(query, nil), do: query
+
+  def search_by_prompt_or_user(query, search_query) do
+    from p in query,
+      left_join: u in User,
+      on: u.id == p.user_id,
+      where: ilike(p.prompt, ^"%#{search_query}%")
+  end
+
+  def order_by(query, nil) do
+    from p in query,
+      order_by: [desc: p.inserted_at]
+  end
+
+  def order_by(query, "desc") do
+    from p in query,
+      join: v in Vote,
+      on: v.prompt_request_id == p.id,
+      group_by: p.id,
+      order_by: [desc: count(v.id)]
+  end
+
+  def order_by(query, "asc") do
+    from p in query,
+      left_join: v in Vote,
+      on: v.prompt_request_id == p.id,
+      group_by: p.id,
+      order_by: [asc: count(v.id)]
+  end
 end

lib/diffuser/generator/prompt_request_result.ex

@@ -1,13 +1,12 @@
 defmodule Diffuser.Generator.PromptRequestResult do
-  use Ecto.Schema
+  use Diffuser.Schema
   use Waffle.Ecto.Schema
   import Ecto.Changeset
   alias Diffuser.Generator.PromptRequest
 
-  @primary_key {:id, Ecto.UUID, autogenerate: true}
   schema "prompt_request_results" do
     field :image, Diffuser.Uploaders.Image.Type
-    belongs_to :prompt_request, PromptRequest, type: :binary_id
+    belongs_to :prompt_request, PromptRequest
 
     timestamps()
   end

lib/diffuser/generator/prompt_request_worker.ex

@@ -5,8 +5,6 @@ defmodule Diffuser.Generator.PromptRequestWorker do
   alias Diffuser.Repo
 
   @path [:code.priv_dir(:diffuser), "python"] |> Path.join()
-  @steps 10
-  @guidance_scale 7.5
 
   def start(%PromptRequest{} = prompt_request) do
     with {:ok, active_prompt} <-
@@ -26,24 +24,38 @@ defmodule Diffuser.Generator.PromptRequestWorker do
     end
   end
 
-  defp update_and_broadcast_progress(%PromptRequest{id: id} = prompt_request, new_status) do
-    {:ok, new_prompt} =
-      Generator.update_prompt_request(prompt_request, %{
-        status: new_status,
-        steps: @steps,
-        guidance_scale: @guidance_scale
+  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, %{id: prompt_id, prompt: prompt}) do
+  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})},
+         ~s(python #{@path}/stable_diffusion.py --prompt "#{prompt}" --output "#{@path}/#{prompt_id}.png" --num-inference-steps #{steps} --guidance-scale #{guidance_scale})},
         [:binary]
       )
 
@@ -56,11 +68,18 @@ defmodule Diffuser.Generator.PromptRequestWorker do
         {:ok, msg}
 
       {^port, {:data, ":step" <> step}} ->
-        Endpoint.broadcast("request:#{prompt_id}", "progress", 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} ->
-        IO.inspect(result, label: "RESULT")
+      {^port, _result} ->
         python_loop(port, prompt_id)
     end
   end