Merge pull request #6 from charlesjlee/fix_typo

fix typo and lint

Author: ezzcodeezzlife

README.md

@@ -4,7 +4,7 @@
 
 - Go to https://labs.openai.com/
 - Open Network Tab in Developer Tools
-- Type a promt and press "Generate"
+- Type a prompt and press "Generate"
 - Look for fetch to https://labs.openai.com/api/labs/tasks
 - In the request header look for authorization then get the Bearer Token
 
@@ -23,8 +23,8 @@ print(generations)
 ```
 
 ```
-✔️  Task created with ID: task-f77yxcsdf3OEm and PROMT: portal to another dimension, digital art
-⌛ Waiting for task to finish .. 
+✔️  Task created with ID: task-f77yxcsdf3OEm and PROMPT: portal to another dimension, digital art
+⌛ Waiting for task to finish...
 🙌 Task completed!
 
 [
@@ -67,8 +67,8 @@ generations = dalle.generate_and_download("portal to another dimension, digital
 ```
 
 ```
-✔️  Task created with ID: task-f77sayxcSGdfOEm and PROMT: portal to another dimension, digital art
-⌛ Waiting for task to finish .. 
+✔️  Task created with ID: task-f77sayxcSGdfOEm and PROMPT: portal to another dimension, digital art
+⌛ Waiting for task to finish...
 🙌 Task completed!
 Download to directory: C:\Users\pc\dalle2
 ✔️  Downloaded:  generation-fAq4Lyxcm7pQVDBQEWJ.jpg
@@ -90,10 +90,10 @@ print(generations)
 ```
 
 ```
-✔️  Task created with ID: task-lm0V4nZasgAFasd7AsStE67 and PROMT: portal to another dimension OVERALL: 1/ 2
-⌛ Waiting for task to finish .. 
+✔️  Task created with ID: task-lm0V4nZasgAFasd7AsStE67 and PROMPT: portal to another dimension OVERALL: 1/ 2
+⌛ Waiting for task to finish...
 ➕  Appended new generations to all_generations
-✔️  Task created with ID: task-WcetZOHt8asdvHb433gi and PROMT: portal to another dimension OVERALL: 2/ 2
+✔️  Task created with ID: task-WcetZOHt8asdvHb433gi and PROMPT: portal to another dimension OVERALL: 2/ 2
 ⌛ Waiting for task to finish .. 
 ➕  Appended new generations to all_generations
 🙌 Task completed!

src/dalle2/dalle2.py

@@ -1,26 +1,25 @@
-from tokenize import String
-import requests
 import json
-import time 
+import os
+import requests
+import time
 import urllib
 import urllib.request
-import os
 
 class Dalle2():
     def __init__(self, bearer):
         self.bearer = bearer
         self.batch_size = 4
 
-    def generate(self, promt):
+    def generate(self, prompt):
         url = "https://labs.openai.com/api/labs/tasks"
         headers = {
             'Authorization': "Bearer " + self.bearer,
-            'Content-Type': "application/json"
+            'Content-Type': "application/json",
         }
         body = {
             "task_type": "text2im",
             "prompt": {
-                "caption": promt,
+                "caption": prompt,
                 "batch_size": self.batch_size,
             }
         }
@@ -30,8 +29,8 @@ def generate(self, promt):
             print(response.text)
             return None
         data = response.json()
-        print("✔️  Task created with ID:", data["id"], "and PROMT:", promt)
-        print("⌛ Waiting for task to finish .. ")
+        print("✔️  Task created with ID:", data["id"], "and PROMPT:", prompt)
+        print("⌛ Waiting for task to finish...")
 
         while True:
             url = "https://labs.openai.com/api/labs/tasks/" + data["id"]
@@ -40,66 +39,61 @@ def generate(self, promt):
             if data["status"] == "succeeded":
                 print("🙌 Task completed!")
                 generations = data["generations"]["data"]
-                return(generations)
-            else:
-                # print("Task not completed yet")
-                time.sleep(3)
-                continue
+                return generations
+
+            time.sleep(3)
 
-    def generate_and_download(self, promt):
-        generations = self.generate(promt)
-        if generations is None:
+    def generate_and_download(self, prompt):
+        generations = self.generate(prompt)
+        if not generations:
             return None
-        
+
         print("Download to directory: " + os.getcwd())
-        #print(generations)
         for generation in generations:
-            imageurl = generation["generation"]["image_path"]
-            id = generation["id"]
+            image_url = generation["generation"]["image_path"]
+            image_id = generation["id"]
+
+            urllib.request.urlretrieve(image_url, image_id +".jpg")
+            print("✔️ Downloaded: ", image_id + ".jpg")
 
-            urllib.request.urlretrieve(imageurl, id +".jpg")
-            print("✔️ Downloaded: " , id + ".jpg")
+        return generations
 
-    def generate_amount(self, promt, amount):
+    def generate_amount(self, prompt, amount):
         url = "https://labs.openai.com/api/labs/tasks"
         headers = {
             'Authorization': "Bearer " + self.bearer,
-            'Content-Type': "application/json"
+            'Content-Type': "application/json",
         }
         body = {
             "task_type": "text2im",
             "prompt": {
-                "caption": promt,
+                "caption": prompt,
                 "batch_size": self.batch_size,
             }
         }
-        
-        
+
         all_generations = []
-        for i in range(1,int(amount / self.batch_size +1)):
+        for i in range(1, int(amount / self.batch_size +1)):
             url = "https://labs.openai.com/api/labs/tasks"
             response = requests.post(url, headers=headers, data=json.dumps(body))
             if response.status_code != 200:
                 print(response.text)
                 return None
             data = response.json()
-            print("✔️ Task created with ID:", data["id"], "and PROMT:", promt, "OVERALL:" , str(i)  + "/" , int(amount / self.batch_size ))
-            print("⌛ Waiting for task to finish .. ")
+            print("✔️ Task created with ID:", data["id"], "and PROMPT:", prompt, "OVERALL:", str(i) + "/", int(amount / self.batch_size))
+            print("⌛ Waiting for task to finish...")
 
             while True:
                 url = "https://labs.openai.com/api/labs/tasks/" + data["id"]
                 response = requests.get(url, headers=headers)
                 data = response.json()
                 if data["status"] == "succeeded":
-                    
                     generations = data["generations"]["data"]
                     print("➕ Appended new generations to all_generations")
                     all_generations.append(generations)
                     break
-                else:
-                    # print("Task not completed yet")
-                    time.sleep(3)
-                    continue
+
+                time.sleep(3)
         print("🙌 Task completed!")
         print(all_generations)
         return all_generations