ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-28 18:21:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adjust dataloader tester to handle videos to test them

Browse Source
This commit is contained in:
Jaret Burkett
2025-10-21 14:47:23 -06:00
parent 0d8a33dc16
commit 5123090f6c
2 changed files with 29 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
testing/test_bucket_dataloader.py
View File

@@ -50,9 +50,13 @@ class FakeAdapter:
class FakeSD:
def __init__(self):
self.adapter = FakeAdapter()
self.use_raw_control_images = False
def encode_control_in_text_embeddings(self, *args, **kwargs):
return None
def get_bucket_divisibility(self):
return 32
dataset_config = DatasetConfig(
dataset_path=dataset_folder,
@@ -120,7 +124,11 @@ for epoch in range(args.epochs):
big_img = img_batch
# big_img = big_img.clamp(-1, 1)
if args.output_path is not None:
save_tensors(big_img, os.path.join(args.output_path, f'{idx}.png'))
if len(img_batch.shape) == 5:
# video
save_tensors(big_img, os.path.join(args.output_path, f'{idx}.webp'), fps=16)
else:
save_tensors(big_img, os.path.join(args.output_path, f'{idx}.png'))
else:
show_tensors(big_img)