add download

Author: ezzcodeezzlife

README.md

@@ -19,10 +19,15 @@ from dalle2 import Dalle2
 dalle = Dalle2("sess-xxxxxxxxxxxxxxxxxxxxxxxxxxxx")
 generations = dalle.generate("portal to another dimension, digital art")
 
+
 print(generations)
 ```
 
 ```
+✔️  Task created with ID: task-f77yxcsdf3OEm and PROMT: portal to another dimension, digital art
+⌛ Waiting for task to finish .. 
+🙌 Task completed!
+
 [
   {
     'id': 'generation-sCnERSYDPP0Zu14fsdXEcKmL',
@@ -52,3 +57,26 @@ print(generations)
 ]
 ```
 
+or download all generations
+
+```python
+from dalle2 import Dalle2
+
+dalle = Dalle2("sess-xxxxxxxxxxxxxxxxxxxxxxxxxxxx")
+generations = dalle.generate_and_download("portal to another dimension, digital art")
+
+```
+
+```
+✔️  Task created with ID: task-f77sayxcSGdfOEm and PROMT: portal to another dimension, digital art
+⌛ Waiting for task to finish .. 
+🙌 Task completed!
+Download to directory: C:\Users\pc\dalle2
+✔️  Downloaded:  generation-fAq4Lyxcm7pQVDBQEWJ.jpg
+✔️  Downloaded:  generation-zqfBC3yyxcPXRlW6zLP.jpg
+✔️  Downloaded:  generation-soR3ryxcoeixzdyHG.jpg
+✔️  Downloaded:  generation-lT5L4yxc2DOiGRwJi.jpg
+✔️  Downloaded:  generation-01DvPryxcq2BX1NOREL.jpg
+✔️  Downloaded:  generation-AAs4yxcczH1vl19FidR.jpg
+
+```

src/dalle2/dalle2.py

@@ -1,6 +1,9 @@
 import requests
 import json
 import time 
+import urllib
+import urllib.request
+import os
 
 class Dalle2():
     def __init__(self, bearer):
@@ -40,4 +43,20 @@ def generate(self, promt):
                 # print("Task not completed yet")
                 time.sleep(3)
                 continue
+
+    def generate_and_download(self, promt):
+        generations = self.generate(promt)
+        if generations is None:
+            return None
+        
+        print("Download to directory: " + os.getcwd())
+        #print(generations)
+        for generation in generations:
+            imageurl = generation["generation"]["image_path"]
+            id = generation["id"]
+
+            urllib.request.urlretrieve(imageurl, id +".jpg")
+            print("✔️  Downloaded: " , id + ".jpg")
+        
+