Mirrors
/
waifu-diffusion
mirror of https://github.com/harubaru/waifu-diffusion.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Final cleanup.

This commit is contained in:
Maw-Fox 2022-11-11 18:09:09 -07:00
parent de221ea42e
commit 6c2d5d8066
1 changed files with 4 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
trainer/diffusers_trainer.py
View File

@ -266,18 +266,15 @@ class ImageStore:
self.validator = Validation( self.validator = Validation(
args.skip_validation, args.skip_validation,
args.extended_validation args.extended_validation
) ).validate
self.resizer = Resize(args.resize, args.data_migration) self.resizer = Resize(args.resize, args.data_migration).resize
self.image_files = [x for x in self.image_files if self.__valid_file(x)] self.image_files = [x for x in self.image_files if self.validator(x)]
def __len__(self) -> int: def __len__(self) -> int:
return len(self.image_files) return len(self.image_files)
def __valid_file(self, f) -> bool:
return self.validator.validate(f)
# iterator returns images as PIL images and their index in the store # iterator returns images as PIL images and their index in the store
def entries_iterator(self) -> Generator[Tuple[Img, int], None, None]: def entries_iterator(self) -> Generator[Tuple[Img, int], None, None]:
for f in range(len(self)): for f in range(len(self)):
@ -285,7 +282,7 @@ class ImageStore:
# get image by index # get image by index
def get_image(self, ref: Tuple[int, int, int]) -> Img: def get_image(self, ref: Tuple[int, int, int]) -> Img:
return self.resizer.resize( return self.resizer(
self.image_files[ref[0]], self.image_files[ref[0]],
ref[1], ref[1],
ref[2] ref[2]