Create dall_e.py

dall_e.py

@@ -0,0 +1,59 @@
+import os, getpass
+
+import numpy as np
+import replicate
+
+
+class DalleImageGenerator:
+    """Dall-e model using Replicate API"""
+    def __init__(self, token=None):
+        if "REPLICATE_API_TOKEN" not in os.environ:
+            if token is not None:
+                os.environ["REPLICATE_API_TOKEN"] = token
+            else:
+                print(f"Please go to https://replicate.com/docs/api for your Replicate API token.")
+                os.environ["REPLICATE_API_TOKEN"] = getpass.getpass(f"Input Replicate API Token:")
+
+        self.dalle = replicate.models.get("kuprel/min-dalle")
+
+    def generate_images(self, text, grid_size, text_adherence=2):
+        urls = self.dalle.predict(text=text, grid_size=grid_size, log2_supercondition_factor=text_adherence)
+        images = get_image(list(urls)[-1])
+        h, w = images.shape[:2]
+        h, w = h // grid_size, w // grid_size
+        return blockshaped(images, h, w)
+
+
+def get_image(url):
+    """download image from a url"""
+    from urllib.request import Request, urlopen
+    import io
+    import PIL.Image as Image
+    hdr = {
+        'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.11 (KHTML, like Gecko) Chrome/23.0.1271.64 Safari/537.11',
+        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
+        'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.3',
+        'Accept-Encoding': 'none',
+        'Accept-Language': 'en-US,en;q=0.8',
+        'Connection': 'keep-alive'}
+
+    # urllib.request.urlretrieve(url, f"local-filename.jpg")
+    req = Request(url, headers=hdr)
+    page = urlopen(req)
+    return np.array(Image.open(io.BytesIO(page.read())))
+
+
+def blockshaped(arr, nrows, ncols):
+    """
+    Return an array of shape (n, nrows, ncols) where
+    n * nrows * ncols = arr.size
+
+    If arr is a 2D array, the returned array should look like n subblocks with
+    each subblock preserving the "physical" layout of arr.
+    """
+    h, w, c = arr.shape
+    assert h % nrows == 0, f"{h} rows is not evenly divisible by {nrows}"
+    assert w % ncols == 0, f"{w} cols is not evenly divisible by {ncols}"
+    return (arr.reshape(h//nrows, nrows, w//ncols, ncols, - 1)
+               .swapaxes(1,2)
+               .reshape(-1, nrows, ncols, 3))