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

Code before:
                    args.class_prompt, text_encoding_pipeline
                )
    validation_embeddings = {}
    if args.validation_prompt is not None:
        if args.remote_text_encoder:
            (validation_embeddings["prompt_embeds"], validation_embeddings["text_ids"]) = (
                compute_remote_text_embeddings(args.validation_prompt)
            )
        else:
            with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
                (validation_embeddings["prompt_embeds"], validation_embeddings["text_ids"]) = compute_text_embeddings(
                    args.validation_prompt, text_encoding_pipeline
                )

    # Init FSDP for text encoder
    if args.fsdp_text_encoder:
        from functools import partial

        from torch.distributed.fsdp import (
            BackwardPrefetch,
            CPUOffload,
            ShardingStrategy,
        )
        from torch.distributed.fsdp import (
            FullyShardedDataParallel as FSDP,
        )
        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

        original_text_encoder = text_encoding_pipeline.text_encoder
        transformer_layer = type(original_text_encoder.model.language_model.layers[0])
        auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={transformer_layer})

        text_encoder_fsdp = FSDP(
            original_text_encoder,
            device_id=accelerator.device,
            sharding_strategy=ShardingStrategy.FULL_SHARD,
            cpu_offload=CPUOffload(offload_params=args.offload),
            auto_wrap_policy=auto_wrap_policy,
            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
            limit_all_gathers=True,
            use_orig_params=True,
        )

        text_encoding_pipeline.text_encoder = text_encoder_fsdp
        dist.barrier()

    # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images),
    # pack the statically computed variables appropriately here. This is so that we don't
    # have to pass them to the dataloader.
    if not train_dataset.custom_instance_prompts:
        prompt_embeds = instance_prompt_hidden_states

Review comment (suggestion):
```suggestion
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
```

Code after:
            if isinstance(prompts, (list, tuple)):
                embeds = [_encode_single(p) for p in prompts]
                prompt_embeds = torch.cat(embeds, dim=0)
            else:
                prompt_embeds = _encode_single(prompts)

            text_ids = Flux2Pipeline._prepare_text_ids(prompt_embeds).to(accelerator.device)
            prompt_embeds = prompt_embeds.to(accelerator.device)
            return prompt_embeds, text_ids

        except Exception as e:
            raise RuntimeError("Remote text encoder inference failed.") from e

    # If no type of tuning is done on the text_encoder and custom instance prompts are NOT
    # provided (i.e. the --instance_prompt is used for all images), we encode the instance prompt once to avoid
    # the redundant encoding.
    if not train_dataset.custom_instance_prompts:
        if args.remote_text_encoder:
            instance_prompt_hidden_states, instance_text_ids = compute_remote_text_embeddings(args.instance_prompt)
        else:
            with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
                instance_prompt_hidden_states, instance_text_ids = compute_text_embeddings(
                    args.instance_prompt, text_encoding_pipeline
                )

    # Handle class prompt for prior-preservation.
    if args.with_prior_preservation:
        if args.remote_text_encoder:
            class_prompt_hidden_states, class_text_ids = compute_remote_text_embeddings(args.class_prompt)
        else:
            with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
                class_prompt_hidden_states, class_text_ids = compute_text_embeddings(
                    args.class_prompt, text_encoding_pipeline
                )
    validation_embeddings = {}
    if args.validation_prompt is not None:
        if args.remote_text_encoder:
            (validation_embeddings["prompt_embeds"], validation_embeddings["text_ids"]) = (
                compute_remote_text_embeddings(args.validation_prompt)
            )
        else:
            with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
                (validation_embeddings["prompt_embeds"], validation_embeddings["text_ids"]) = compute_text_embeddings(
                    args.validation_prompt, text_encoding_pipeline
                )

    # Init FSDP for text encoder
    if args.fsdp_text_encoder:
        fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator)
        text_encoder_fsdp = wrap_with_fsdp(
            model=text_encoding_pipeline.text_encoder,