Pop args not allowed (#15)

jackalcooper · Mar 19, 2024 · a8bbcfb · a8bbcfb
1 parent 8ff5fc0
commit a8bbcfb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/lib/live_xl_web/live/prompt.ex b/lib/live_xl_web/live/prompt.ex
@@ -80,6 +80,10 @@ defmodule LiveXLWeb.PromptLive do
 
     args =
       args
+      |> Map.reject(fn
+        {k, _} when is_bitstring(k) -> String.starts_with?(k, "_")
+        _ -> false
+      end)
       |> Map.put("saved_image", saved_image)
       |> Map.put("seed", seed)
       |> Map.put("num_inference_steps", 2)

diff --git a/priv/runner.py b/priv/runner.py
@@ -100,11 +100,12 @@ def recv_loop_jsonl(stream):
             try:
                 args = message["payload"]["args"]
                 seed = args.pop("seed", 1)
+                saved_image = args.pop("saved_image")
                 args["generator"] = torch.Generator(device="cuda").manual_seed(seed)
                 start_time = time.time()
                 images = lightning.pipe(**args).images
                 end_time = time.time()
-                images[0].save(args["saved_image"])
+                images[0].save(saved_image)
             except Exception as e:
                 send_jsonl(
                     {