Initial Commit

1 files changed, 64 insertions, 0 deletions

diff --git a/dalle_mini_api.nim b/dalle_mini_api.nim
new file mode 100644
index 0000000..d11d8f1
--- /dev/null
+++ b/dalle_mini_api.nim
@@ -0,0 +1,64 @@
+import argparse
+import std/httpclient
+import std/base64
+import std/oids
+import std/json
+import std/strutils
+
+var parser = newParser:
+    help("Abuse the dalle-mini API from the cli!")
+    option("-o", "--output", help="Write output to this folder", default=some("/tmp/"))
+    option("-n", "--number", help="number of images you'd like", default=some("9"))
+    arg("prompt", nargs = -1)
+
+
+try:
+    let arguments = parser.parse()
+    var client = newHttpClient(userAgent="Mozilla/5.0 (Windows NT 10.0; rv:91.0) Gecko/20100101 Firefox/91.0")
+    client.headers = newHttpHeaders({
+     "Accept":"application/json",
+     "Accept-Language":"en-US,en;q=0.5",
+     "Accept-Encoding":"gzip, deflate, br",
+     "Referer":"https://hf.space/",
+     "Content-Type":"application/json",
+     "Origin":"https://hf.space",
+     "DNT":"1",
+     "Connection":"keep-alive",
+     "Sec-Fetch-Dest":"empty",
+     "Sec-Fetch-Mode":"cors",
+     "Sec-Fetch-Site":"cross-site"
+        })
+
+    for parsedPrompt in arguments.prompt:
+        echo "Now generating: " & parsedPrompt
+        var data = "{\"prompt\":\"" & parsedPrompt & "\"}"
+        var response: Response
+        var max = (parseInt(arguments.number)/9 + 1)
+        for i in countup(0,toInt(max)):
+            echo "New run..."
+
+            echo "Waiting for response..."
+            response = client.request("https://bf.dallemini.ai/generate", httpMethod = HttpPost, body=data)
+            while response.status != "200 OK":
+                echo "Could not get response due to " & response.status & "! Trying again in 15 seconds!"
+                sleep 15
+                response = client.request("https://bf.dallemini.ai/generate", httpMethod = HttpPost, body=data)
+    
+            var JSONifiedResponse = parseJSON(response.body)
+            for encodedImage in JSONifiedResponse["images"]:
+                var stringImage = $encodedImage
+                #var image = base64.decode(stringImage.replace("\\").replace("\"").replace("="))
+                #echo stringImage.replace("\\n").replace("\"")
+                var image = base64.decode(stringImage.replace("\\n").replace("\""))
+                var filename = arguments.output & parsedPrompt & " - " & $oids.genOid() & ".jpg"
+                echo "writing image to " & filename
+                writeFile(filename, image)
+
+
+except ShortCircuit as e:
+  if e.flag == "argparse_help":
+    echo parser.help
+    quit(1)
+except UsageError:
+  stderr.writeLine getCurrentExceptionMsg()
+  quit(1)