diff options
Diffstat (limited to 'dallegen.nim')
-rw-r--r-- | dallegen.nim | 109 |
1 files changed, 109 insertions, 0 deletions
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) |