improved the number parameter and added an option to upscale

4 files changed, 119 insertions, 71 deletions

diff --git a/README.md b/README.md
index ed0950e..d462772 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,6 @@
-<h1>Dalle Mini API Abuser</h1>
+<h1>Dallegen : A Dalle Mini API Abuser</h1>
 <p>Allows you to use <a href=https://huggingface.co/spaces/flax-community/dalle-mini>Dalle Mini</a> from the CLI via the public API.</p>
 <h2>Usage</h2>
-Abuse the dalle-mini API from the cli!
-
 Usage:
    [options] [prompt ...]
 
@@ -12,4 +10,6 @@ Arguments:
 Options:
   -h, --help
   -o, --output=OUTPUT        Write output to this folder (default: /tmp/)
-  -n, --number=NUMBER        number of images you'd like (default: 9)
+  -n, --number=NUMBER        the number of times you'd like to request for each prompt. Expect 9 per request. (default: 1)
+  -u, --upscale              Use SwiniR to upscale the image
+
diff --git a/dalle_mini_api.nim b/dalle_mini_api.nim
deleted file mode 100644
index d11d8f1..0000000
--- a/dalle_mini_api.nim
+++ /dev/null
@@ -1,64 +0,0 @@
-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)
diff --git a/dallegen.nim b/dallegen.nim
new file mode 100644
index 0000000..f032e2c
--- /dev/null
+++ b/dallegen.nim
@@ -0,0 +1,109 @@
+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="the number of times you'd like to request for each prompt. Expect 9 per request.", default=some("1"))
+    flag("-u", "--upscale", help="Use SwiniR to upscale the image")
+    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)-1)
+        for i in countup(0,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 "DALLE-Mini: Could not get response due to " & response.status & "! Trying again in 15 seconds!"
+                sleep 15000
+                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
+                
+                if arguments.upscale:
+                    var encoded = stringImage.replace("\\n").replace("\"")
+                    var data = %*{
+                        "action": "predict",
+                        "data": [
+                            "data:image/png;base64," & encoded
+                        ],
+                        "session_hash": "faq0rxk7k3"
+                    }
+
+                    echo "Starting to upscale!"
+                    var response = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/push/", httpMethod = HttpPost, body = $data)
+                    while response.status != "200 OK":
+                        echo "SwinIR: Could not get response due to " & response.status & "! Trying again in 15 seconds!"
+                        sleep 15000
+                        response = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/push/", httpMethod = HttpPost, body = $data)
+                    
+                    var hash = parseJSON(response.body)
+                    delete(hash, "queue_position")
+                    
+                    
+                    var isItDoneYet = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/status/", httpMethod = HttpPost, body = $hash)
+                    var doneStatus = parseJSON(isItDoneYet.body)
+                    var waitingList = parseJSON(response.body)
+                    echo "we are currently in position number " &  $waitingList["queue_position"] & " in the queue!"
+                    
+                    while $doneStatus["status"] == "\"QUEUED\"" or $doneStatus["status"] == "\"PENDING\"":
+                        echo "Waiting for 3 more seconds"
+                        sleep 3000
+                        isItDoneYet = client.request("https://hf.space/embed/akhaliq/SwinIR/api/queue/status/", httpMethod = HttpPost, body = $hash)
+                        doneStatus = parseJSON(isItDoneYet.body)
+                    
+                    if $doneStatus["status"] == "\"FAILED\"":
+                        echo "ERROR WHILE GETTING STATUS:"
+                        echo $doneStatus
+                        quit 1
+
+                    echo "Done!"
+                    var outputdata = $doneStatus["data"]["data"][0]
+                    var image = base64.decode(outputdata.replace("\"")[22..^1])
+                    var filename = arguments.output & parsedPrompt & " - " & $oids.genOid() & ".png"
+                    echo "writing image to " & filename
+                    writeFile(filename, image)
+
+                else:
+                    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)
diff --git a/manpage b/manpage
index 2455d61..dd222ed 100644
--- a/manpage
+++ b/manpage
@@ -1,7 +1,7 @@
 .\" DO NOT MODIFY THIS FILE!  It was generated by help2man 1.49.2.
-.TH DALLE_MINI_API "Michael" "June 2022" "dalle_mini_api 1.0" "User Commands"
+.TH DALLEGEN "1" "June 2022" "dallegen 2" "User Commands"
 .SH NAME
-dalle_mini_api \- Dalle Mini API ABuser
+dallegen \- manual page for dallegen 2
 .SH DESCRIPTION
 Abuse the dalle\-mini API from the cli!
 .SS "Usage:"
@@ -18,4 +18,7 @@ Abuse the dalle\-mini API from the cli!
 Write output to this folder (default: /tmp/)
 .TP
 \fB\-n\fR, \fB\-\-number\fR=\fI\,NUMBER\/\fR
-number of images you'd like (default: 9)
+the number of times you'd like to request for each prompt. Expect 9 per request. (default: 1)
+.TP
+\fB\-u\fR, \fB\-\-upscale\fR
+Use SwiniR to upscale the image