Skip to content

Commit 173e1b1

Browse files
[Examples] Uniform notations in train_flux_lora (#10011)
[Examples] uniform naming notations since the in parameter `size` represents `args.resolution`, I thus replace the `args.resolution` inside DreamBoothData with `size`. And revise some notations such as `center_crop`. Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
1 parent e46e139 commit 173e1b1

File tree

1 file changed

+5
-6
lines changed

1 file changed

+5
-6
lines changed

examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -971,6 +971,7 @@ class DreamBoothDataset(Dataset):
971971

972972
def __init__(
973973
self,
974+
args,
974975
instance_data_root,
975976
instance_prompt,
976977
class_prompt,
@@ -980,10 +981,8 @@ def __init__(
980981
class_num=None,
981982
size=1024,
982983
repeats=1,
983-
center_crop=False,
984984
):
985985
self.size = size
986-
self.center_crop = center_crop
987986

988987
self.instance_prompt = instance_prompt
989988
self.custom_instance_prompts = None
@@ -1075,11 +1074,11 @@ def __init__(
10751074
# flip
10761075
image = train_flip(image)
10771076
if args.center_crop:
1078-
y1 = max(0, int(round((image.height - args.resolution) / 2.0)))
1079-
x1 = max(0, int(round((image.width - args.resolution) / 2.0)))
1077+
y1 = max(0, int(round((image.height - self.size) / 2.0)))
1078+
x1 = max(0, int(round((image.width - self.size) / 2.0)))
10801079
image = train_crop(image)
10811080
else:
1082-
y1, x1, h, w = train_crop.get_params(image, (args.resolution, args.resolution))
1081+
y1, x1, h, w = train_crop.get_params(image, (self.size, self.size))
10831082
image = crop(image, y1, x1, h, w)
10841083
image = train_transforms(image)
10851084
self.pixel_values.append(image)
@@ -1827,6 +1826,7 @@ def load_model_hook(models, input_dir):
18271826

18281827
# Dataset and DataLoaders creation:
18291828
train_dataset = DreamBoothDataset(
1829+
args=args,
18301830
instance_data_root=args.instance_data_dir,
18311831
instance_prompt=args.instance_prompt,
18321832
train_text_encoder_ti=args.train_text_encoder_ti,
@@ -1836,7 +1836,6 @@ def load_model_hook(models, input_dir):
18361836
class_num=args.num_class_images,
18371837
size=args.resolution,
18381838
repeats=args.repeats,
1839-
center_crop=args.center_crop,
18401839
)
18411840

18421841
train_dataloader = torch.utils.data.DataLoader(

0 commit comments

Comments
 (0)