Mirrors
/
EveryDream2trainer
mirror of https://github.com/victorchall/EveryDream2trainer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Always convert YAML tags to strings

This commit is contained in:
Augusto de la Torre 2023-04-23 00:36:32 +02:00
parent 9b6c01a7f0
commit 2e8056174d
3 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
data/dataset.py
View File

@ -36,7 +36,7 @@ class Tag:
return Tag(data) return Tag(data)
if isinstance(data, dict): if isinstance(data, dict):
value = data.get("tag") value = str(data.get("tag"))
weight = data.get("weight") weight = data.get("weight")
if value: if value:
return Tag(value, weight) return Tag(value, weight)

2
test/test_data_resolver.py
View File

@ -131,7 +131,7 @@ class TestResolve(unittest.TestCase):
self.assertEqual(len(items), 6) self.assertEqual(len(items), 6)
self.assertEqual(set(image_paths), set([IMAGE_1_PATH, IMAGE_2_PATH, IMAGE_3_PATH] * 2)) self.assertEqual(set(image_paths), set([IMAGE_1_PATH, IMAGE_2_PATH, IMAGE_3_PATH] * 2))
self.assertEqual(set(captions), {}'caption for test1', 'test2', 'test3', 'caption for test1', 'caption for test2', 'test3'}) self.assertEqual(set(captions), {'caption for test1', 'test2', 'test3', 'caption for test1', 'caption for test2', 'test3'})
undersized_images = list(filter(lambda i: i.is_undersized, items)) undersized_images = list(filter(lambda i: i.is_undersized, items))
self.assertEqual(len(undersized_images), 2) self.assertEqual(len(undersized_images), 2)

3
test/test_dataset.py
View File

@ -77,6 +77,7 @@ class TestDataset(TestCase):
- tag: from .yml - tag: from .yml
- tag: with weight - tag: with weight
weight: 0.5 weight: 0.5
- tag: 1234.5
""")) """))
actual = Dataset.from_path(".").image_configs actual = Dataset.from_path(".").image_configs
@ -94,7 +95,7 @@ class TestDataset(TestCase):
rating=1.1, rating=1.1,
max_caption_length=1024, max_caption_length=1024,
main_prompts="A complex caption", main_prompts="A complex caption",
tags= { Tag("from .yml"), Tag("with weight", weight=0.5) } tags= { Tag("from .yml"), Tag("with weight", weight=0.5), Tag("1234.5") }
) )
} }
self.assertEqual(expected, actual) self.assertEqual(expected, actual)