stable-generate: use new api

1 files changed, 16 insertions, 48 deletions

diff --git a/krebs/5pkgs/simple/stable-generate/default.nix b/krebs/5pkgs/simple/stable-generate/default.nix
index fac26161..dc9c826f 100644
--- a/krebs/5pkgs/simple/stable-generate/default.nix
+++ b/krebs/5pkgs/simple/stable-generate/default.nix
@@ -1,64 +1,32 @@
 { pkgs, lib, ... }:
 
-pkgs.writers.writeDashBin "stable-generate" ''
+pkgs.writers.writeBashBin "stable-generate" ''
   set -efu
 
   export PATH=${lib.makeBinPath [
+    pkgs.coreutils
     pkgs.curl
     pkgs.jq
   ]}
 
   STABLE_URL=''${STABLE_URL:-http://stable-confusion.r}
 
-  PAYLOAD=$(jq -cn --arg query "$*" '{fn_index: 51, data: [
-    $query,
-    "",
-    "None",
-    "None",
-    20, # sampling steps
-    "Euler a", # sampling method
-    false, # restore faces
-    false,
-    1,
-    1,
-    7,
-    -1,
-    -1,
-    0,
-    0,
-    0,
-    false,
-    512, #probably resolution
-    512, #probably resolution
-    false,
-    0.7,
-    0,
-    0,
-    "None",
-    "",
-    false,
-    false,
-    false,
-    "",
-    "Seed",
-    "",
-    "Nothing",
-    "",
-    true,
-    false,
-    false,
-    null,
-    "",
-  ""], session_hash: "hello_this_is_dog"}')
+  PAYLOAD=$(jq -cn --arg prompt "$*" '{
+    prompt: $prompt
+  }')
 
-  data=$(curl -Ssf "$STABLE_URL/run/predict/" \
+  filename=$(mktemp)
+  curl -Ssf "$STABLE_URL/sdapi/v1/txt2img" \
     -X POST \
     --Header 'Content-Type: application/json' \
-    --data "$PAYLOAD"
-  )
-  export data
+    --data "$PAYLOAD" |
+      jq -r '.images[0]' |
+      base64 --decode > "$filename"
 
-  filename=$(jq -rn 'env.data | fromjson.data[0][0].name')
-
-  echo "$STABLE_URL/file=$filename"
+  if test -t 1; then
+    echo "$filename"
+  else
+    cat "$filename"
+    rm "$filename"
+  fi
 ''