From bdd404f79b1bb0f9c53eafb67990269d22e33fe3 Mon Sep 17 00:00:00 2001
From: Michael Chung <michael.chung@databites.ca>
Date: Tue, 14 Nov 2023 00:53:55 -0500
Subject: [PATCH] Works with models folder dynamically loading models from the
 models folder

---
 animation/Rerender_A_Video/sd_model_cfg.py | 44 ++++++++++++++++++----
 animation/Rerender_A_Video/webUI.py        |  3 +-
 2 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/animation/Rerender_A_Video/sd_model_cfg.py b/animation/Rerender_A_Video/sd_model_cfg.py
index 4fb7af0..7c38c20 100644
--- a/animation/Rerender_A_Video/sd_model_cfg.py
+++ b/animation/Rerender_A_Video/sd_model_cfg.py
@@ -1,14 +1,42 @@
 # The model dict is used for webUI only
+from typing import Dict
+import unittest
+import pathlib
 
-model_dict = {
-    'Stable Diffusion 1.5': '',
-    'revAnimated_v11': 'models/revAnimated_v11.safetensors',
-    'realisticVisionV20_v20': 'models/realisticVisionV20_v20.safetensors',
-    'CounterfeitV30_25': 'models/CounterfeitV30_25.safetensors',
-    'ghibli_style_offset': 'models/ghibli_style_offset.safetensors'
-}
+
+def path_to_dict(path: pathlib.Path) -> Dict[str, str]:
+    path = pathlib.Path(path)
+    result = {}
+    for file in path.iterdir():
+        if file.is_file() and file.suffix == '.safetensors':
+            # Convert file name to camel case with underscores between words
+            name = '_'.join(word.capitalize() for word in file.stem.split('_'))
+            # Add entry to dictionary
+            result[name] = str(file.relative_to(path.parent))
+    result['Stable Diffusion 1.5'] = ''
+    print("result: ", result)
+    return result
+
+
+# model_dict = {
+#     'Stable Diffusion 1.5': '',
+#     'revAnimated_v11': 'models/revAnimated_v11.safetensors',
+#     'realisticVisionV20_v20': 'models/realisticVisionV20_v20.safetensors',
+#     'CounterfeitV30_25': 'models/CounterfeitV30_25.safetensors',
+#     'ghibli_style_offset': 'models/ghibli_style_offset.safetensors'
+# }
+model_dict = path_to_dict("./models")
 
 loRA_dict = {
     'None': '',
     'Dark Sushi': 'loRAs/dark_sushi'
-}
\ No newline at end of file
+}
+
+class TestPathToDict(unittest.TestCase):
+    def test_path_to_dict(self):
+        # Create a temporary directory
+        result = path_to_dict(pathlib.Path("/home/tensor/code/storyteller/storyteller-ml/animation/Rerender_A_Video/models"))
+        print(result)
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/animation/Rerender_A_Video/webUI.py b/animation/Rerender_A_Video/webUI.py
index 88a09e7..8171a7f 100644
--- a/animation/Rerender_A_Video/webUI.py
+++ b/animation/Rerender_A_Video/webUI.py
@@ -32,7 +32,7 @@ from src.freeu import freeu_forward
 from src.img_util import find_flat_region, numpy2tensor
 from src.video_util import (frame_to_video, get_fps, get_frame_count,
                             prepare_frames)
-
+from pathlib import Path
 inversed_model_dict = dict()
 for k, v in model_dict.items():
     inversed_model_dict[v] = k
@@ -41,6 +41,7 @@ to_tensor = T.PILToTensor()
 blur = T.GaussianBlur(kernel_size=(9, 9), sigma=(18, 18))
 
 
+
 class ProcessingState(Enum):
     NULL = 0
     FIRST_IMG = 1