Works with models folder dynamically loading models from the models folder

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'
-}
+}
+
+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()

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