Update diffusers_trainer.py

This commit is contained in:
laksjdjf 2022-11-21 13:57:26 +09:00 committed by GitHub
parent 511ee9e6d2
commit 5787f7d080
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions

View File

@ -719,7 +719,7 @@ def main():
# load dataset # load dataset
store = ImageStore(args.dataset) store = ImageStore(args.dataset)
dataset = AspectDataset(store, tokenizer) dataset = AspectDataset(store, tokenizer, ucg=args.ucg)
bucket = AspectBucket(store, args.num_buckets, args.batch_size, args.bucket_side_min, args.bucket_side_max, 64, args.resolution * args.resolution, 2.0) bucket = AspectBucket(store, args.num_buckets, args.batch_size, args.bucket_side_min, args.bucket_side_max, 64, args.resolution * args.resolution, 2.0)
sampler = AspectBucketSampler(bucket=bucket, num_replicas=world_size, rank=rank) sampler = AspectBucketSampler(bucket=bucket, num_replicas=world_size, rank=rank)