Repository: huggingface/diffusers
File: examples/dreambooth/train_dreambooth_lora_flux2.py (Python)

Code before:
    if args.gradient_checkpointing:
        transformer.enable_gradient_checkpointing()

    if args.lora_layers is not None:
        target_modules = [layer.strip() for layer in args.lora_layers.split(",")]
    else:
        target_modules = ["to_k", "to_q", "to_v", "to_out.0"]

    # now we will add new LoRA weights the transformer layers
    transformer_lora_config = LoraConfig(
        r=args.rank,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        init_lora_weights="gaussian",
        target_modules=target_modules,
    )
    transformer.add_adapter(transformer_lora_config)

    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
    def save_model_hook(models, weights, output_dir):
        transformer_lora_layers_to_save = None
        modules_to_save = {}

        if is_fsdp:
            for model in models:
                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
                    state_dict = accelerator.get_state_dict(models)

                    if accelerator.is_main_process:
                        transformer_lora_layers_to_save = get_peft_model_state_dict(
                            unwrap_model(model),
                            state_dict=state_dict,
                        )
                        transformer_lora_layers_to_save = {
                            k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v
                            for k, v in transformer_lora_layers_to_save.items()
                        }
                        modules_to_save["transformer"] = model

                        # make sure to pop weight so that corresponding model is not saved again
                        if weights:
                            weights.pop()
        else:
            if accelerator.is_main_process:
                transformer_lora_layers_to_save = None
                modules_to_save = {}

Review comment (refactor):
Let's simplify this block of code a bit:

```py
transformer_cls = type(unwrap_model(transformer))

def _to_cpu_contiguous(sd):
    return {
        k: (v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v)
        for k, v in sd.items()
    }

# 1) Validate and pick the transformer model
modules_to_save: dict[str, Any] = {}
transformer_model = None

for m in models:
    if isinstance(unwrap_model(m), transformer_cls):
        transformer_model = m
        modules_to_save["transformer"] = m
    else:
        raise ValueError(f"unexpected save model: {m.__class__}")

if transformer_model is None:
    raise ValueError("No transformer model found in `models`.")

# 2) Optionally gather FSDP state dict once
state_dict = accelerator.get_state_dict(models) if is_fsdp else None

# 3) Only main process materializes the LoRA state dict
transformer_lora_layers_to_save = None
if accelerator.is_main_process:
    peft_kwargs = {}
    if is_fsdp:
        peft_kwargs["state_dict"] = state_dict

    transformer_lora_layers_to_save = get_peft_model_state_dict(
        unwrap_model(transformer_model) if is_fsdp else transformer_model,
        **peft_kwargs,
    )

    if is_fsdp:
        transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save)

        # make sure to pop weight so that corresponding model is not saved again
        if weights:
            weights.pop()
```

We can move `_to_cpu_contiguous()` to the `training_utils.py` module.

Code after:
            scheduler=None,
            revision=args.revision,
        )

    if args.gradient_checkpointing:
        transformer.enable_gradient_checkpointing()

    if args.lora_layers is not None:
        target_modules = [layer.strip() for layer in args.lora_layers.split(",")]
    else:
        target_modules = ["to_k", "to_q", "to_v", "to_out.0"]

    # now we will add new LoRA weights the transformer layers
    transformer_lora_config = LoraConfig(
        r=args.rank,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        init_lora_weights="gaussian",
        target_modules=target_modules,
    )
    transformer.add_adapter(transformer_lora_config)

    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
    def save_model_hook(models, weights, output_dir):
        transformer_cls = type(unwrap_model(transformer))

        # 1) Validate and pick the transformer model
        modules_to_save: dict[str, Any] = {}
        transformer_model = None

        for model in models:
            if isinstance(unwrap_model(model), transformer_cls):
                transformer_model = model
                modules_to_save["transformer"] = model
            else:
                raise ValueError(f"unexpected save model: {model.__class__}")

        if transformer_model is None:
            raise ValueError("No transformer model found in 'models'")

        # 2) Optionally gather FSDP state dict once
        state_dict = accelerator.get_state_dict(model) if is_fsdp else None

        # 3) Only main process materializes the LoRA state dict
        transformer_lora_layers_to_save = None
        if accelerator.is_main_process: