mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
108 lines
2.6 KiB
Python
108 lines
2.6 KiB
Python
import time
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from torchvision import transforms
|
|
import sys
|
|
import os
|
|
import cv2
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from toolkit.paths import SD_SCRIPTS_ROOT
|
|
|
|
from toolkit.image_utils import show_img
|
|
|
|
sys.path.append(SD_SCRIPTS_ROOT)
|
|
|
|
from library.model_util import load_vae
|
|
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
|
|
from toolkit.data_loader import AiToolkitDataset, get_dataloader_from_datasets, \
|
|
trigger_dataloader_setup_epoch
|
|
from toolkit.config_modules import DatasetConfig
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('dataset_folder', type=str, default='input')
|
|
parser.add_argument('--epochs', type=int, default=1)
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
dataset_folder = args.dataset_folder
|
|
resolution = 1024
|
|
bucket_tolerance = 64
|
|
batch_size = 1
|
|
|
|
|
|
##
|
|
|
|
dataset_config = DatasetConfig(
|
|
dataset_path=dataset_folder,
|
|
resolution=resolution,
|
|
caption_ext='json',
|
|
default_caption='default',
|
|
buckets=True,
|
|
bucket_tolerance=bucket_tolerance,
|
|
poi='person',
|
|
augmentations=[
|
|
{
|
|
'method': 'RandomBrightnessContrast',
|
|
'brightness_limit': (-0.3, 0.3),
|
|
'contrast_limit': (-0.3, 0.3),
|
|
'brightness_by_max': False,
|
|
'p': 1.0
|
|
},
|
|
{
|
|
'method': 'HueSaturationValue',
|
|
'hue_shift_limit': (-0, 0),
|
|
'sat_shift_limit': (-40, 40),
|
|
'val_shift_limit': (-40, 40),
|
|
'p': 1.0
|
|
},
|
|
# {
|
|
# 'method': 'RGBShift',
|
|
# 'r_shift_limit': (-20, 20),
|
|
# 'g_shift_limit': (-20, 20),
|
|
# 'b_shift_limit': (-20, 20),
|
|
# 'p': 1.0
|
|
# },
|
|
]
|
|
|
|
|
|
)
|
|
|
|
dataloader: DataLoader = get_dataloader_from_datasets([dataset_config], batch_size=batch_size)
|
|
|
|
|
|
# run through an epoch ang check sizes
|
|
dataloader_iterator = iter(dataloader)
|
|
for epoch in range(args.epochs):
|
|
for batch in dataloader:
|
|
batch: 'DataLoaderBatchDTO'
|
|
img_batch = batch.tensor
|
|
|
|
chunks = torch.chunk(img_batch, batch_size, dim=0)
|
|
# put them so they are size by side
|
|
big_img = torch.cat(chunks, dim=3)
|
|
big_img = big_img.squeeze(0)
|
|
|
|
min_val = big_img.min()
|
|
max_val = big_img.max()
|
|
|
|
big_img = (big_img / 2 + 0.5).clamp(0, 1)
|
|
|
|
# convert to image
|
|
img = transforms.ToPILImage()(big_img)
|
|
|
|
show_img(img)
|
|
|
|
time.sleep(1.0)
|
|
# if not last epoch
|
|
if epoch < args.epochs - 1:
|
|
trigger_dataloader_setup_epoch(dataloader)
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
print('done')
|