fix: eval schedule string (#125)

Author: hanxiao

discoart/helper.py

@@ -692,17 +692,18 @@ def _version_check(package: str = None, github_repo: str = None):
 
 def _eval_scheduling_str(val) -> List[float]:
     if isinstance(val, str):
-        r = eval(val)
-    elif isinstance(val, (int, float, bool)):
-        r = [val] * _MAX_DIFFUSION_STEPS
+        val = eval(val)
+
+    if isinstance(val, (int, float, bool)):
+        val = [val] * _MAX_DIFFUSION_STEPS
+    elif isinstance(val, (list, tuple)):
+        if len(val) != _MAX_DIFFUSION_STEPS:
+            raise ValueError(
+                f'invalid scheduling string: {val} the schedule steps should be exactly {_MAX_DIFFUSION_STEPS}'
+            )
     else:
         raise ValueError(f'unsupported scheduling type: {val}: {type(val)}')
-
-    if len(r) != _MAX_DIFFUSION_STEPS:
-        raise ValueError(
-            f'invalid scheduling string: {val} the schedule steps should be exactly {_MAX_DIFFUSION_STEPS}'
-        )
-    return r
+    return val
 
 
 def _get_current_schedule(schedule_table: Dict, t: int) -> 'SimpleNamespace':

tests/test_config.py

@@ -1,4 +1,5 @@
 from discoart.config import default_args, save_config, load_config, export_python
+from discoart.helper import _eval_scheduling_str
 
 
 def test_export_load_config(tmpfile):
@@ -23,3 +24,13 @@ def test_format_config():
 
 def test_export_python():
     assert export_python(default_args)
+
+
+def test_eval_schedule_string():
+    assert _eval_scheduling_str('1') == [1] * 1000
+    assert _eval_scheduling_str('[1] * 1000') == [1] * 1000
+    assert _eval_scheduling_str(1) == [1] * 1000
+    assert _eval_scheduling_str('1.') == [1] * 1000
+    assert _eval_scheduling_str('True') == [True] * 1000
+    assert _eval_scheduling_str('False') == [False] * 1000
+    assert _eval_scheduling_str(True) == [True] * 1000