adding the argument to load_dataset function call

Author: YanivDorGalron

examples/controlnet/train_controlnet_flax.py

@@ -521,6 +521,7 @@ def make_train_dataset(args, tokenizer, batch_size=None):
             args.dataset_config_name,
             cache_dir=args.cache_dir,
             streaming=args.streaming,
+            trust_remote_code=args.trust_remote_code
         )
     else:
         if args.train_data_dir is not None:
@@ -532,6 +533,7 @@ def make_train_dataset(args, tokenizer, batch_size=None):
                 dataset = load_dataset(
                     args.train_data_dir,
                     cache_dir=args.cache_dir,
+                    trust_remote_code=args.trust_remote_code
                 )
         # See more about loading custom images at
         # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script

examples/controlnet/train_controlnet_flux.py

@@ -13,30 +13,29 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 
+import accelerate
 import argparse
 import copy
 import functools
 import logging
 import math
+import numpy as np
 import os
 import random
 import shutil
-from contextlib import nullcontext
-from pathlib import Path
-
-import accelerate
-import numpy as np
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import transformers
+from PIL import Image
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import DistributedType, ProjectConfiguration, set_seed
+from contextlib import nullcontext
 from datasets import load_dataset
 from huggingface_hub import create_repo, upload_folder
 from packaging import version
-from PIL import Image
+from pathlib import Path
 from torchvision import transforms
 from tqdm.auto import tqdm
 from transformers import (
@@ -60,7 +59,6 @@
 from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available
 from diffusers.utils.torch_utils import is_compiled_module
 
-
 if is_wandb_available():
     import wandb
 
@@ -73,7 +71,7 @@
 
 
 def log_validation(
-    vae, flux_transformer, flux_controlnet, args, accelerator, weight_dtype, step, is_final_validation=False
+        vae, flux_transformer, flux_controlnet, args, accelerator, weight_dtype, step, is_final_validation=False
 ):
     logger.info("Running validation... ")
 
@@ -266,7 +264,7 @@ def parse_args(input_args=None):
         type=str,
         default=None,
         help="Path to pretrained controlnet model or model identifier from huggingface.co/models."
-        " If not specified controlnet weights are initialized from unet.",
+             " If not specified controlnet weights are initialized from unet.",
     )
     parser.add_argument(
         "--variant",
@@ -668,11 +666,11 @@ def parse_args(input_args=None):
         raise ValueError("`--validation_prompt` must be set if `--validation_image` is set")
 
     if (
-        args.validation_image is not None
-        and args.validation_prompt is not None
-        and len(args.validation_image) != 1
-        and len(args.validation_prompt) != 1
-        and len(args.validation_image) != len(args.validation_prompt)
+            args.validation_image is not None
+            and args.validation_prompt is not None
+            and len(args.validation_image) != 1
+            and len(args.validation_prompt) != 1
+            and len(args.validation_image) != len(args.validation_prompt)
     ):
         raise ValueError(
             "Must provide either 1 `--validation_image`, 1 `--validation_prompt`,"
@@ -695,10 +693,12 @@ def get_train_dataset(args, accelerator):
             args.dataset_name,
             args.dataset_config_name,
             cache_dir=args.cache_dir,
+            trust_remote_code=args.trust_remote_code
         )
     if args.jsonl_for_train is not None:
         # load from json
-        dataset = load_dataset("json", data_files=args.jsonl_for_train, cache_dir=args.cache_dir)
+        dataset = load_dataset("json", data_files=args.jsonl_for_train, cache_dir=args.cache_dir,
+                               trust_remote_code=args.trust_remote_code)
         dataset = dataset.flatten_indices()
     # Preprocessing the datasets.
     # We need to tokenize inputs and targets.
@@ -1018,7 +1018,7 @@ def load_model_hook(models, input_dir):
 
     if args.scale_lr:
         args.learning_rate = (
-            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+                args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
         )
 
     # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
@@ -1130,7 +1130,7 @@ def compute_embeddings(batch, proportion_empty_prompts, flux_controlnet_pipeline
         len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
         num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
         num_training_steps_for_scheduler = (
-            args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
+                args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
         )
     else:
         num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes

examples/controlnet/train_controlnet_sd3.py

@@ -650,12 +650,14 @@ def make_train_dataset(args, tokenizer_one, tokenizer_two, tokenizer_three, acce
             args.dataset_name,
             args.dataset_config_name,
             cache_dir=args.cache_dir,
+            trust_remote_code=args.trust_remote_code
         )
     else:
         if args.train_data_dir is not None:
             dataset = load_dataset(
                 args.train_data_dir,
                 cache_dir=args.cache_dir,
+                trust_remote_code=args.trust_remote_code
             )
         # See more about loading custom images at
         # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script

examples/controlnet/train_controlnet_sdxl.py

@@ -645,12 +645,14 @@ def get_train_dataset(args, accelerator):
             args.dataset_config_name,
             cache_dir=args.cache_dir,
             data_dir=args.train_data_dir,
+            trust_remote_code=args.trust_remote_code
         )
     else:
         if args.train_data_dir is not None:
             dataset = load_dataset(
                 args.train_data_dir,
                 cache_dir=args.cache_dir,
+                trust_remote_code=args.trust_remote_code
             )
         # See more about loading custom images at
         # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script

examples/research_projects/pixart/train_pixart_controlnet_hf.py

@@ -771,6 +771,7 @@ def load_model_hook(models, input_dir):
             args.dataset_config_name,
             cache_dir=args.cache_dir,
             data_dir=args.train_data_dir,
+            trust_remote_code=args.trust_remote_code
         )
     else:
         data_files = {}
@@ -780,6 +781,7 @@ def load_model_hook(models, input_dir):
             "imagefolder",
             data_files=data_files,
             cache_dir=args.cache_dir,
+            trust_remote_code=args.trust_remote_code
         )
         # See more about loading custom images at
         # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder

examples/t2i_adapter/train_t2i_adapter_sdxl.py

@@ -637,12 +637,14 @@ def get_train_dataset(args, accelerator):
             args.dataset_name,
             args.dataset_config_name,
             cache_dir=args.cache_dir,
+            trust_remote_code=args.trust_remote_code
         )
     else:
         if args.train_data_dir is not None:
             dataset = load_dataset(
                 args.train_data_dir,
                 cache_dir=args.cache_dir,
+                trust_remote_code=args.trust_remote_code
             )
         # See more about loading custom images at
         # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script