From 7ac6b3a46cb73604aeff865f61b5d6dcab72fa4e Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Fri, 2 Jun 2023 12:48:39 -0700
Subject: [PATCH] some more options

---
 TPDNE_utils/tpdne.py | 33 +++++++++++++++++++++++++++------
 setup.py             |  2 +-
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/TPDNE_utils/tpdne.py b/TPDNE_utils/tpdne.py
index fc0ff8a..7291059 100644
--- a/TPDNE_utils/tpdne.py
+++ b/TPDNE_utils/tpdne.py
@@ -14,12 +14,24 @@ def sample_image_and_save_repeatedly(
     output_path: str = './out/random',     # path to the output image, without extension (will be saved as webp)
     *,
     call_every_ms: int = 250,              # how often to sample
-    tmp_dir = '/tmp',                      # to store temporary images, before symbolically linking to the output path
-    num_rotated_tmp_images = 10,
-    verbose: bool = True
+    tmp_dir: str = '/tmp',                 # to store temporary images, before symbolically linking to the output path
+    num_rotated_tmp_images: int = 10,
+    image_format: str = 'webp',
+    verbose: bool = True,
+    quality = 99,
+    generate_favicon: bool = True,
+    favicon_size: int = 32
 ):
+    assert 0 < quality <= 100
+    assert favicon_size in {16, 32}
+    assert image_format in {'webp', 'png', 'jpeg'}
+
     tmp_dir = Path(tmp_dir)
-    output_path = Path(output_path).with_suffix('.webp')
+    output_path = Path(output_path)
+    assert output_path.suffix == '', 'output path suffix will be automatically determined by `image_format` keyword arg'
+
+    output_path = output_path.with_suffix(f'.{image_format}')
+
     call_every_seconds = call_every_ms / 1000
 
     assert tmp_dir.is_dir()
@@ -32,11 +44,20 @@ def sample_image_and_save_repeatedly(
         image_tensor = fn()
 
         tmp_image_index = (tmp_image_index + 1) % num_rotated_tmp_images
-        tmp_path = str(tmp_dir / f'{tmp_image_index}.webp')
+        tmp_path = str(tmp_dir / f'{tmp_image_index}.{image_format}')
 
-        Image.fromarray(image_tensor, 'RGB').save(tmp_path, format = 'webp')
+        pil_image = Image.fromarray(image_tensor, 'RGB')
+        pil_image.save(tmp_path, format = image_format, quality = quality)
         os.system(f'ln -nfs {tmp_path} {output_path}')
 
+        if generate_favicon:
+            tmp_favicon_path = str(tmp_dir / f'favicon_{tmp_image_index}.png')
+            output_favicon_path = output_path.parents[0] / 'favicon.png'
+
+            small_pil_image = pil_image.resize((favicon_size, favicon_size))
+            small_pil_image.save(tmp_favicon_path)
+            os.system(f'ln -nfs {tmp_favicon_path} {output_favicon_path}')
+
         elapsed = time() - start
 
         if verbose:
diff --git a/setup.py b/setup.py
index e869b12..722045a 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'TPDNE-utils',
   packages = find_packages(exclude=[]),
-  version = '0.0.2',
+  version = '0.0.3',
   license='MIT',
   description = 'TPDNE',
   author = 'Phil Wang',