Unverified Commit ee1516e5 authored by Evan Han's avatar Evan Han Committed by GitHub
Browse files

[train_dreambooth_lora_lumina2] Add LANCZOS as the default interpolation mode...

[train_dreambooth_lora_lumina2] Add LANCZOS as the default interpolation mode for image resizing (#11491)

[ADD] interpolation
parent ec932399
......@@ -599,6 +599,15 @@ def parse_args(input_args=None):
"Defaults to precision dtype used for training to save memory"
),
)
parser.add_argument(
"--image_interpolation_mode",
type=str,
default="lanczos",
choices=[
f.lower() for f in dir(transforms.InterpolationMode) if not f.startswith("__") and not f.endswith("__")
],
help="The image interpolation method to use for resizing images.",
)
parser.add_argument(
"--offload",
action="store_true",
......@@ -724,7 +733,11 @@ class DreamBoothDataset(Dataset):
self.instance_images.extend(itertools.repeat(img, repeats))
self.pixel_values = []
train_resize = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR)
interpolation = getattr(transforms.InterpolationMode, args.image_interpolation_mode.upper(), None)
if interpolation is None:
raise ValueError(f"Unsupported interpolation mode: {args.image_interpolation_mode}")
train_resize = transforms.Resize(size, interpolation=interpolation)
train_crop = transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size)
train_flip = transforms.RandomHorizontalFlip(p=1.0)
train_transforms = transforms.Compose(
......@@ -768,7 +781,7 @@ class DreamBoothDataset(Dataset):
self.image_transforms = transforms.Compose(
[
transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
transforms.Resize(size, interpolation=interpolation),
transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment