train_control_lora_flux.py

#!/usr/bin/env python
# coding=utf-8
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and

importargparse
importcopy
importlogging
importmath
importos
importrandom
importshutil
fromcontextlibimportnullcontext
frompathlibimportPath

importaccelerate
importnumpyasnp
importtorch
importtransformers
fromaccelerateimportAccelerator
fromaccelerate.loggingimportget_logger
fromaccelerate.utilsimportDistributedType, ProjectConfiguration, set_seed
fromdatasetsimportload_dataset
fromhuggingface_hubimportcreate_repo, upload_folder
frompackagingimportversion
frompeftimportLoraConfig, set_peft_model_state_dict
frompeft.utilsimportget_peft_model_state_dict
fromPILimportImage
fromtorchvisionimporttransforms
fromtqdm.autoimporttqdm

importdiffusers
fromdiffusersimportAutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxControlPipeline, FluxTransformer2DModel
fromdiffusers.optimizationimportget_scheduler
fromdiffusers.training_utilsimport (
cast_training_params,
compute_density_for_timestep_sampling,
compute_loss_weighting_for_sd3,
free_memory,
)
fromdiffusers.utilsimportcheck_min_version, is_wandb_available, load_image, make_image_grid
fromdiffusers.utils.hub_utilsimportload_or_create_model_card, populate_model_card
fromdiffusers.utils.torch_utilsimportis_compiled_module


ifis_wandb_available():
importwandb

# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
check_min_version("0.34.0.dev0")

logger=get_logger(__name__)

NORM_LAYER_PREFIXES= ["norm_q", "norm_k", "norm_added_q", "norm_added_k"]


defencode_images(pixels: torch.Tensor, vae: torch.nn.Module, weight_dtype):
pixel_latents=vae.encode(pixels.to(vae.dtype)).latent_dist.sample()
pixel_latents= (pixel_latents-vae.config.shift_factor) *vae.config.scaling_factor
returnpixel_latents.to(weight_dtype)


deflog_validation(flux_transformer, args, accelerator, weight_dtype, step, is_final_validation=False):
logger.info("Running validation... ")

ifnotis_final_validation:
flux_transformer=accelerator.unwrap_model(flux_transformer)
pipeline=FluxControlPipeline.from_pretrained(
args.pretrained_model_name_or_path,
transformer=flux_transformer,
torch_dtype=weight_dtype,
 )
else:
transformer=FluxTransformer2DModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="transformer", torch_dtype=weight_dtype
 )
initial_channels=transformer.config.in_channels
pipeline=FluxControlPipeline.from_pretrained(
args.pretrained_model_name_or_path,
transformer=transformer,
torch_dtype=weight_dtype,
 )
pipeline.load_lora_weights(args.output_dir)
assertpipeline.transformer.config.in_channels==initial_channels*2, (
f"{pipeline.transformer.config.in_channels=}"
 )

pipeline.to(accelerator.device)
pipeline.set_progress_bar_config(disable=True)

ifargs.seedisNone:
generator=None
else:
generator=torch.Generator(device=accelerator.device).manual_seed(args.seed)

iflen(args.validation_image) ==len(args.validation_prompt):
validation_images=args.validation_image
validation_prompts=args.validation_prompt
eliflen(args.validation_image) ==1:
validation_images=args.validation_image*len(args.validation_prompt)
validation_prompts=args.validation_prompt
eliflen(args.validation_prompt) ==1:
validation_images=args.validation_image
validation_prompts=args.validation_prompt*len(args.validation_image)
else:
raiseValueError(
"number of `args.validation_image` and `args.validation_prompt` should be checked in `parse_args`"
 )

image_logs= []
ifis_final_validationortorch.backends.mps.is_available():
autocast_ctx=nullcontext()
else:
autocast_ctx=torch.autocast(accelerator.device.type, weight_dtype)

forvalidation_prompt, validation_imageinzip(validation_prompts, validation_images):
validation_image=load_image(validation_image)
# maybe need to inference on 1024 to get a good image
validation_image=validation_image.resize((args.resolution, args.resolution))

images= []

for_inrange(args.num_validation_images):
withautocast_ctx:
image=pipeline(
prompt=validation_prompt,
control_image=validation_image,
num_inference_steps=50,
guidance_scale=args.guidance_scale,
generator=generator,
max_sequence_length=512,
height=args.resolution,
width=args.resolution,
 ).images[0]
image=image.resize((args.resolution, args.resolution))
images.append(image)
image_logs.append(
 {"validation_image": validation_image, "images": images, "validation_prompt": validation_prompt}
 )

tracker_key="test"ifis_final_validationelse"validation"
fortrackerinaccelerator.trackers:
iftracker.name=="tensorboard":
forloginimage_logs:
images=log["images"]
validation_prompt=log["validation_prompt"]
validation_image=log["validation_image"]
formatted_images= []
formatted_images.append(np.asarray(validation_image))
forimageinimages:
formatted_images.append(np.asarray(image))
formatted_images=np.stack(formatted_images)
tracker.writer.add_images(validation_prompt, formatted_images, step, dataformats="NHWC")

eliftracker.name=="wandb":
formatted_images= []
forloginimage_logs:
images=log["images"]
validation_prompt=log["validation_prompt"]
validation_image=log["validation_image"]
formatted_images.append(wandb.Image(validation_image, caption="Conditioning"))
forimageinimages:
image=wandb.Image(image, caption=validation_prompt)
formatted_images.append(image)

tracker.log({tracker_key: formatted_images})
else:
logger.warning(f"image logging not implemented for {tracker.name}")

delpipeline
free_memory()
returnimage_logs


defsave_model_card(repo_id: str, image_logs=None, base_model=str, repo_folder=None):
img_str=""
ifimage_logsisnotNone:
img_str="You can find some example images below.\n\n"
fori, loginenumerate(image_logs):
images=log["images"]
validation_prompt=log["validation_prompt"]
validation_image=log["validation_image"]
validation_image.save(os.path.join(repo_folder, "image_control.png"))
img_str+=f"prompt: {validation_prompt}\n"
images= [validation_image] +images
make_image_grid(images, 1, len(images)).save(os.path.join(repo_folder, f"images_{i}.png"))
img_str+=f"![images_{i})](./images_{i}.png)\n"

model_description=f"""
# control-lora-{repo_id}

These are Control LoRA weights trained on {base_model} with new type of conditioning.
{img_str}

## License

Please adhere to the licensing terms as described [here](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)
"""

model_card=load_or_create_model_card(
repo_id_or_path=repo_id,
from_training=True,
license="other",
base_model=base_model,
model_description=model_description,
inference=True,
 )

tags= [
"flux",
"flux-diffusers",
"text-to-image",
"diffusers",
"control-lora",
"diffusers-training",
"lora",
 ]
model_card=populate_model_card(model_card, tags=tags)

model_card.save(os.path.join(repo_folder, "README.md"))


defparse_args(input_args=None):
parser=argparse.ArgumentParser(description="Simple example of a Control LoRA training script.")
parser.add_argument(
"--pretrained_model_name_or_path",
type=str,
default=None,
required=True,
help="Path to pretrained model or model identifier from huggingface.co/models.",
 )
parser.add_argument(
"--variant",
type=str,
default=None,
help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
 )
parser.add_argument(
"--revision",
type=str,
default=None,
required=False,
help="Revision of pretrained model identifier from huggingface.co/models.",
 )
parser.add_argument(
"--output_dir",
type=str,
default="control-lora",
help="The output directory where the model predictions and checkpoints will be written.",
 )
parser.add_argument(
"--cache_dir",
type=str,
default=None,
help="The directory where the downloaded models and datasets will be stored.",
 )
parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
parser.add_argument(
"--resolution",
type=int,
default=1024,
help=(
"The resolution for input images, all the images in the train/validation dataset will be resized to this"
" resolution"
 ),
 )
parser.add_argument(
"--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
 )
parser.add_argument("--num_train_epochs", type=int, default=1)
parser.add_argument(
"--max_train_steps",
type=int,
default=None,
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
 )
parser.add_argument(
"--checkpointing_steps",
type=int,
default=500,
help=(
"Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
"In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
"Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
"See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
"instructions."
 ),
 )
parser.add_argument(
"--checkpoints_total_limit",
type=int,
default=None,
help=("Max number of checkpoints to store."),
 )
parser.add_argument(
"--resume_from_checkpoint",
type=str,
default=None,
help=(
"Whether training should be resumed from a previous checkpoint. Use a path saved by"
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
 ),
 )
parser.add_argument(
"--proportion_empty_prompts",
type=float,
default=0,
help="Proportion of image prompts to be replaced with empty strings. Defaults to 0 (no prompt replacement).",
 )
parser.add_argument(
"--rank",
type=int,
default=4,
help=("The dimension of the LoRA update matrices."),
 )
parser.add_argument("--use_lora_bias", action="store_true", help="If training the bias of lora_B layers.")
parser.add_argument(
"--lora_layers",
type=str,
default=None,
help=(
'The transformer modules to apply LoRA training on. Please specify the layers in a comma seperated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only'
 ),
 )
parser.add_argument(
"--gaussian_init_lora",
action="store_true",
help="If using the Gaussian init strategy. When False, we follow the original LoRA init strategy.",
 )
parser.add_argument("--train_norm_layers", action="store_true", help="Whether to train the norm scales.")
parser.add_argument(
"--gradient_accumulation_steps",
type=int,
default=1,
help="Number of updates steps to accumulate before performing a backward/update pass.",
 )
parser.add_argument(
"--gradient_checkpointing",
action="store_true",
help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
 )
parser.add_argument(
"--learning_rate",
type=float,
default=5e-6,
help="Initial learning rate (after the potential warmup period) to use.",
 )
parser.add_argument(
"--scale_lr",
action="store_true",
default=False,
help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
 )
parser.add_argument(
"--lr_scheduler",
type=str,
default="constant",
help=(
'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
' "constant", "constant_with_warmup"]'
 ),
 )
parser.add_argument(
"--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
 )
parser.add_argument(
"--lr_num_cycles",
type=int,
default=1,
help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
 )
parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.")
parser.add_argument(
"--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
 )

parser.add_argument(
"--dataloader_num_workers",
type=int,
default=0,
help=(
"Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
 ),
 )
parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
parser.add_argument(
"--hub_model_id",
type=str,
default=None,
help="The name of the repository to keep in sync with the local `output_dir`.",
 )
parser.add_argument(
"--logging_dir",
type=str,
default="logs",
help=(
"[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
 ),
 )
parser.add_argument(
"--allow_tf32",
action="store_true",
help=(
"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
 ),
 )
parser.add_argument(
"--report_to",
type=str,
default="tensorboard",
help=(
'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
 ),
 )
parser.add_argument(
"--mixed_precision",
type=str,
default=None,
choices=["no", "fp16", "bf16"],
help=(
"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"
" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
 ),
 )
parser.add_argument(
"--dataset_name",
type=str,
default=None,
help=(
"The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"
" dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
" or to a folder containing files that 🤗 Datasets can understand."
 ),
 )
parser.add_argument(
"--dataset_config_name",
type=str,
default=None,
help="The config of the Dataset, leave as None if there's only one config.",
 )
parser.add_argument(
"--image_column", type=str, default="image", help="The column of the dataset containing the target image."
 )
parser.add_argument(
"--conditioning_image_column",
type=str,
default="conditioning_image",
help="The column of the dataset containing the control conditioning image.",
 )
parser.add_argument(
"--caption_column",
type=str,
default="text",
help="The column of the dataset containing a caption or a list of captions.",
 )
parser.add_argument("--log_dataset_samples", action="store_true", help="Whether to log somple dataset samples.")
parser.add_argument(
"--max_train_samples",
type=int,
default=None,
help=(
"For debugging purposes or quicker training, truncate the number of training examples to this "
"value if set."
 ),
 )
parser.add_argument(
"--validation_prompt",
type=str,
default=None,
nargs="+",
help=(
"A set of prompts evaluated every `--validation_steps` and logged to `--report_to`."
" Provide either a matching number of `--validation_image`s, a single `--validation_image`"
" to be used with all prompts, or a single prompt that will be used with all `--validation_image`s."
 ),
 )
parser.add_argument(
"--validation_image",
type=str,
default=None,
nargs="+",
help=(
"A set of paths to the control conditioning image be evaluated every `--validation_steps`"
" and logged to `--report_to`. Provide either a matching number of `--validation_prompt`s, a"
" a single `--validation_prompt` to be used with all `--validation_image`s, or a single"
" `--validation_image` that will be used with all `--validation_prompt`s."
 ),
 )
parser.add_argument(
"--num_validation_images",
type=int,
default=1,
help="Number of images to be generated for each `--validation_image`, `--validation_prompt` pair",
 )
parser.add_argument(
"--validation_steps",
type=int,
default=100,
help=(
"Run validation every X steps. Validation consists of running the prompt"
" `args.validation_prompt` multiple times: `args.num_validation_images`"
" and logging the images."
 ),
 )
parser.add_argument(
"--tracker_project_name",
type=str,
default="flux_train_control_lora",
help=(
"The `project_name` argument passed to Accelerator.init_trackers for"
" more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"
 ),
 )
parser.add_argument(
"--jsonl_for_train",
type=str,
default=None,
help="Path to the jsonl file containing the training data.",
 )

parser.add_argument(
"--guidance_scale",
type=float,
default=30.0,
help="the guidance scale used for transformer.",
 )

parser.add_argument(
"--upcast_before_saving",
action="store_true",
help=(
"Whether to upcast the trained transformer layers to float32 before saving (at the end of training). "
"Defaults to precision dtype used for training to save memory"
 ),
 )

parser.add_argument(
"--weighting_scheme",
type=str,
default="none",
choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"],
help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'),
 )
parser.add_argument(
"--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
 )
parser.add_argument(
"--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
 )
parser.add_argument(
"--mode_scale",
type=float,
default=1.29,
help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
 )
parser.add_argument(
"--offload",
action="store_true",
help="Whether to offload the VAE and the text encoders to CPU when they are not used.",
 )

ifinput_argsisnotNone:
args=parser.parse_args(input_args)
else:
args=parser.parse_args()

ifargs.dataset_nameisNoneandargs.jsonl_for_trainisNone:
raiseValueError("Specify either `--dataset_name` or `--jsonl_for_train`")

ifargs.dataset_nameisnotNoneandargs.jsonl_for_trainisnotNone:
raiseValueError("Specify only one of `--dataset_name` or `--jsonl_for_train`")

ifargs.proportion_empty_prompts<0orargs.proportion_empty_prompts>1:
raiseValueError("`--proportion_empty_prompts` must be in the range [0, 1].")

ifargs.validation_promptisnotNoneandargs.validation_imageisNone:
raiseValueError("`--validation_image` must be set if `--validation_prompt` is set")

ifargs.validation_promptisNoneandargs.validation_imageisnotNone:
raiseValueError("`--validation_prompt` must be set if `--validation_image` is set")

if (
args.validation_imageisnotNone
andargs.validation_promptisnotNone
andlen(args.validation_image) !=1
andlen(args.validation_prompt) !=1
andlen(args.validation_image) !=len(args.validation_prompt)
 ):
raiseValueError(
"Must provide either 1 `--validation_image`, 1 `--validation_prompt`,"
" or the same number of `--validation_prompt`s and `--validation_image`s"
 )

ifargs.resolution%8!=0:
raiseValueError(
"`--resolution` must be divisible by 8 for consistently sized encoded images between the VAE and the Flux transformer."
 )

returnargs


defget_train_dataset(args, accelerator):
dataset=None
ifargs.dataset_nameisnotNone:
# Downloading and loading a dataset from the hub.
dataset=load_dataset(
args.dataset_name,
args.dataset_config_name,
cache_dir=args.cache_dir,
 )
ifargs.jsonl_for_trainisnotNone:
# load from json
dataset=load_dataset("json", data_files=args.jsonl_for_train, cache_dir=args.cache_dir)
dataset=dataset.flatten_indices()
# Preprocessing the datasets.
# We need to tokenize inputs and targets.
column_names=dataset["train"].column_names

# 6. Get the column names for input/target.
ifargs.image_columnisNone:
image_column=column_names[0]
logger.info(f"image column defaulting to {image_column}")
else:
image_column=args.image_column
ifimage_columnnotincolumn_names:
raiseValueError(
f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
 )

ifargs.caption_columnisNone:
caption_column=column_names[1]
logger.info(f"caption column defaulting to {caption_column}")
else:
caption_column=args.caption_column
ifcaption_columnnotincolumn_names:
raiseValueError(
f"`--caption_column` value '{args.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
 )

ifargs.conditioning_image_columnisNone:
conditioning_image_column=column_names[2]
logger.info(f"conditioning image column defaulting to {conditioning_image_column}")
else:
conditioning_image_column=args.conditioning_image_column
ifconditioning_image_columnnotincolumn_names:
raiseValueError(
f"`--conditioning_image_column` value '{args.conditioning_image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
 )

withaccelerator.main_process_first():
train_dataset=dataset["train"].shuffle(seed=args.seed)
ifargs.max_train_samplesisnotNone:
train_dataset=train_dataset.select(range(args.max_train_samples))
returntrain_dataset


defprepare_train_dataset(dataset, accelerator):
image_transforms=transforms.Compose(
 [
transforms.Resize((args.resolution, args.resolution), interpolation=transforms.InterpolationMode.BILINEAR),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
 ]
 )

defpreprocess_train(examples):
images= [
 (image.convert("RGB") ifnotisinstance(image, str) elseImage.open(image).convert("RGB"))
forimageinexamples[args.image_column]
 ]
images= [image_transforms(image) forimageinimages]

conditioning_images= [
 (image.convert("RGB") ifnotisinstance(image, str) elseImage.open(image).convert("RGB"))
forimageinexamples[args.conditioning_image_column]
 ]
conditioning_images= [image_transforms(image) forimageinconditioning_images]
examples["pixel_values"] =images
examples["conditioning_pixel_values"] =conditioning_images

is_caption_list=isinstance(examples[args.caption_column][0], list)
ifis_caption_list:
examples["captions"] = [max(example, key=len) forexampleinexamples[args.caption_column]]
else:
examples["captions"] =list(examples[args.caption_column])

returnexamples

withaccelerator.main_process_first():
dataset=dataset.with_transform(preprocess_train)

returndataset


defcollate_fn(examples):
pixel_values=torch.stack([example["pixel_values"] forexampleinexamples])
pixel_values=pixel_values.to(memory_format=torch.contiguous_format).float()
conditioning_pixel_values=torch.stack([example["conditioning_pixel_values"] forexampleinexamples])
conditioning_pixel_values=conditioning_pixel_values.to(memory_format=torch.contiguous_format).float()
captions= [example["captions"] forexampleinexamples]
return {"pixel_values": pixel_values, "conditioning_pixel_values": conditioning_pixel_values, "captions": captions}


defmain(args):
ifargs.report_to=="wandb"andargs.hub_tokenisnotNone:
raiseValueError(
"You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token."
" Please use `huggingface-cli login` to authenticate with the Hub."
 )
ifargs.use_lora_biasandargs.gaussian_init_lora:
raiseValueError("`gaussian` LoRA init scheme isn't supported when `use_lora_bias` is True.")

logging_out_dir=Path(args.output_dir, args.logging_dir)

iftorch.backends.mps.is_available() andargs.mixed_precision=="bf16":
# due to pytorch#99272, MPS does not yet support bfloat16.
raiseValueError(
"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
 )

accelerator_project_config=ProjectConfiguration(project_dir=args.output_dir, logging_dir=str(logging_out_dir))

accelerator=Accelerator(
gradient_accumulation_steps=args.gradient_accumulation_steps,
mixed_precision=args.mixed_precision,
log_with=args.report_to,
project_config=accelerator_project_config,
 )

# Disable AMP for MPS. A technique for accelerating machine learning computations on iOS and macOS devices.
iftorch.backends.mps.is_available():
logger.info("MPS is enabled. Disabling AMP.")
accelerator.native_amp=False

# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
# DEBUG, INFO, WARNING, ERROR, CRITICAL
level=logging.INFO,
 )
logger.info(accelerator.state, main_process_only=False)

ifaccelerator.is_local_main_process:
transformers.utils.logging.set_verbosity_warning()
diffusers.utils.logging.set_verbosity_info()
else:
transformers.utils.logging.set_verbosity_error()
diffusers.utils.logging.set_verbosity_error()

# If passed along, set the training seed now.
ifargs.seedisnotNone:
set_seed(args.seed)

# Handle the repository creation
ifaccelerator.is_main_process:
ifargs.output_dirisnotNone:
os.makedirs(args.output_dir, exist_ok=True)

ifargs.push_to_hub:
repo_id=create_repo(
repo_id=args.hub_model_idorPath(args.output_dir).name, exist_ok=True, token=args.hub_token
 ).repo_id

# Load models. We will load the text encoders later in a pipeline to compute
# embeddings.
vae=AutoencoderKL.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="vae",
revision=args.revision,
variant=args.variant,
 )
vae_scale_factor=2** (len(vae.config.block_out_channels) -1)
flux_transformer=FluxTransformer2DModel.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="transformer",
revision=args.revision,
variant=args.variant,
 )
logger.info("All models loaded successfully")

noise_scheduler=FlowMatchEulerDiscreteScheduler.from_pretrained(
args.pretrained_model_name_or_path,
subfolder="scheduler",
 )
noise_scheduler_copy=copy.deepcopy(noise_scheduler)
vae.requires_grad_(False)
flux_transformer.requires_grad_(False)

# cast down and move to the CPU
weight_dtype=torch.float32
ifaccelerator.mixed_precision=="fp16":
weight_dtype=torch.float16
elifaccelerator.mixed_precision=="bf16":
weight_dtype=torch.bfloat16

# let's not move the VAE to the GPU yet.
vae.to(dtype=torch.float32) # keep the VAE in float32.
flux_transformer.to(dtype=weight_dtype, device=accelerator.device)

# enable image inputs
withtorch.no_grad():
initial_input_channels=flux_transformer.config.in_channels
new_linear=torch.nn.Linear(
flux_transformer.x_embedder.in_features*2,
flux_transformer.x_embedder.out_features,
bias=flux_transformer.x_embedder.biasisnotNone,
dtype=flux_transformer.dtype,
device=flux_transformer.device,
 )
new_linear.weight.zero_()
new_linear.weight[:, :initial_input_channels].copy_(flux_transformer.x_embedder.weight)
ifflux_transformer.x_embedder.biasisnotNone:
new_linear.bias.copy_(flux_transformer.x_embedder.bias)
flux_transformer.x_embedder=new_linear

asserttorch.all(flux_transformer.x_embedder.weight[:, initial_input_channels:].data==0)
flux_transformer.register_to_config(in_channels=initial_input_channels*2, out_channels=initial_input_channels)

ifargs.train_norm_layers:
forname, paraminflux_transformer.named_parameters():
ifany(kinnameforkinNORM_LAYER_PREFIXES):
param.requires_grad=True

ifargs.lora_layersisnotNone:
ifargs.lora_layers!="all-linear":
target_modules= [layer.strip() forlayerinargs.lora_layers.split(",")]
# add the input layer to the mix.
if"x_embedder"notintarget_modules:
target_modules.append("x_embedder")
elifargs.lora_layers=="all-linear":
target_modules=set()
forname, moduleinflux_transformer.named_modules():
ifisinstance(module, torch.nn.Linear):
target_modules.add(name)
target_modules=list(target_modules)
else:
target_modules= [
"x_embedder",
"attn.to_k",
"attn.to_q",
"attn.to_v",
"attn.to_out.0",
"attn.add_k_proj",
"attn.add_q_proj",
"attn.add_v_proj",
"attn.to_add_out",
"ff.net.0.proj",
"ff.net.2",
"ff_context.net.0.proj",
"ff_context.net.2",
 ]
transformer_lora_config=LoraConfig(
r=args.rank,
lora_alpha=args.rank,
init_lora_weights="gaussian"ifargs.gaussian_init_loraelseTrue,
target_modules=target_modules,
lora_bias=args.use_lora_bias,
 )
flux_transformer.add_adapter(transformer_lora_config)

defunwrap_model(model):
model=accelerator.unwrap_model(model)
model=model._orig_modifis_compiled_module(model) elsemodel
returnmodel

# `accelerate` 0.16.0 will have better support for customized saving
ifversion.parse(accelerate.__version__) >=version.parse("0.16.0"):

defsave_model_hook(models, weights, output_dir):
ifaccelerator.is_main_process:
transformer_lora_layers_to_save=None

formodelinmodels:
ifisinstance(unwrap_model(model), type(unwrap_model(flux_transformer))):
model=unwrap_model(model)
transformer_lora_layers_to_save=get_peft_model_state_dict(model)
ifargs.train_norm_layers:
transformer_norm_layers_to_save= {
f"transformer.{name}": param
forname, paraminmodel.named_parameters()
ifany(kinnameforkinNORM_LAYER_PREFIXES)
 }
transformer_lora_layers_to_save= {
**transformer_lora_layers_to_save,
**transformer_norm_layers_to_save,
 }
else:
raiseValueError(f"unexpected save model: {model.__class__}")

# make sure to pop weight so that corresponding model is not saved again
ifweights:
weights.pop()

FluxControlPipeline.save_lora_weights(
output_dir,
transformer_lora_layers=transformer_lora_layers_to_save,
 )

defload_model_hook(models, input_dir):
transformer_=None

ifnotaccelerator.distributed_type==DistributedType.DEEPSPEED:
whilelen(models) >0:
model=models.pop()

ifisinstance(model, type(unwrap_model(flux_transformer))):
transformer_=model
else:
raiseValueError(f"unexpected save model: {model.__class__}")
else:
transformer_=FluxTransformer2DModel.from_pretrained(
args.pretrained_model_name_or_path, subfolder="transformer"
 ).to(accelerator.device, weight_dtype)

# Handle input dimension doubling before adding adapter
withtorch.no_grad():
initial_input_channels=transformer_.config.in_channels
new_linear=torch.nn.Linear(
transformer_.x_embedder.in_features*2,
transformer_.x_embedder.out_features,
bias=transformer_.x_embedder.biasisnotNone,
dtype=transformer_.dtype,
device=transformer_.device,
 )
new_linear.weight.zero_()
new_linear.weight[:, :initial_input_channels].copy_(transformer_.x_embedder.weight)
iftransformer_.x_embedder.biasisnotNone:
new_linear.bias.copy_(transformer_.x_embedder.bias)
transformer_.x_embedder=new_linear
transformer_.register_to_config(in_channels=initial_input_channels*2)

transformer_.add_adapter(transformer_lora_config)

lora_state_dict=FluxControlPipeline.lora_state_dict(input_dir)
transformer_lora_state_dict= {
f"{k.replace('transformer.', '')}": v
fork, vinlora_state_dict.items()
ifk.startswith("transformer.") and"lora"ink
 }
incompatible_keys=set_peft_model_state_dict(
transformer_, transformer_lora_state_dict, adapter_name="default"
 )
ifincompatible_keysisnotNone:
# check only for unexpected keys
unexpected_keys=getattr(incompatible_keys, "unexpected_keys", None)
ifunexpected_keys:
logger.warning(
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
f" {unexpected_keys}. "
 )
ifargs.train_norm_layers:
transformer_norm_state_dict= {
k: v
fork, vinlora_state_dict.items()
ifk.startswith("transformer.") andany(norm_kinkfornorm_kinNORM_LAYER_PREFIXES)
 }
transformer_._transformer_norm_layers=FluxControlPipeline._load_norm_into_transformer(
transformer_norm_state_dict,
transformer=transformer_,
discard_original_layers=False,
 )

# Make sure the trainable params are in float32. This is again needed since the base models
# are in `weight_dtype`. More details:
# https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
ifargs.mixed_precision=="fp16":
models= [transformer_]
# only upcast trainable parameters (LoRA) into fp32
cast_training_params(models)

accelerator.register_save_state_pre_hook(save_model_hook)
accelerator.register_load_state_pre_hook(load_model_hook)

# Make sure the trainable params are in float32.
ifargs.mixed_precision=="fp16":
models= [flux_transformer]
# only upcast trainable parameters (LoRA) into fp32
cast_training_params(models, dtype=torch.float32)