jhj0517 commited on
Commit
289ab0d
·
1 Parent(s): a81a186

Add yaml utils

Browse files
Files changed (1) hide show
  1. modules/utils/files_manager.py +34 -1
modules/utils/files_manager.py CHANGED
@@ -1,8 +1,41 @@
1
  import os
2
  import fnmatch
3
-
4
  from gradio.utils import NamedString
5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
 
7
  def get_media_files(folder_path, include_sub_directory=False):
8
  video_extensions = ['*.mp4', '*.mkv', '*.flv', '*.avi', '*.mov', '*.wmv']
 
1
  import os
2
  import fnmatch
3
+ import yaml
4
  from gradio.utils import NamedString
5
 
6
+ from modules.utils.paths import DEFAULT_PARAMETERS_CONFIG_PATH
7
+
8
+
9
+ class YAMLDumper(yaml.SafeDumper):
10
+ def __init__(self, *args, **kwargs):
11
+ super().__init__(*args, **kwargs)
12
+ self.add_representer(str, self.str_presenter)
13
+
14
+ def write_line_break(self, data=None):
15
+ super().write_line_break(data)
16
+ if len(self.indents) == 1:
17
+ super().write_line_break()
18
+
19
+ @staticmethod
20
+ def str_presenter(dumper, data):
21
+ special_chars = set(' \n\t:-[]{},&*#?|>!%@`"\'')
22
+ if any(char in special_chars for char in data) or data == '':
23
+ return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='"')
24
+ return dumper.represent_scalar('tag:yaml.org,2002:str', data)
25
+
26
+
27
+ def load_yaml(path: str = DEFAULT_PARAMETERS_CONFIG_PATH):
28
+ with open(path, 'r', encoding='utf-8') as file:
29
+ config = yaml.safe_load(file)
30
+ return config
31
+
32
+
33
+ def save_yaml(data: dict, path: str = DEFAULT_PARAMETERS_CONFIG_PATH):
34
+ data = yaml.dump(data, sort_keys=False, default_flow_style=False, Dumper=YAMLDumper, allow_unicode=True)
35
+ with open(path, 'w', encoding='utf-8') as file:
36
+ file.write(data)
37
+ return path
38
+
39
 
40
  def get_media_files(folder_path, include_sub_directory=False):
41
  video_extensions = ['*.mp4', '*.mkv', '*.flv', '*.avi', '*.mov', '*.wmv']