|  | 
| 1 | 1 | # Training Control LoRA with Flux | 
| 2 | 2 | 
 | 
| 3 |  | -This example shows how train Control LoRA with Flux to condition it with additional structural controls (like depth maps, poses, etc.).  | 
|  | 3 | +This (experimental) example shows how train Control LoRA with [Flux](https://huggingface.co/black-forest-labs/FLUX.1-dev) to condition it with additional structural controls (like depth maps, poses, etc.).  | 
| 4 | 4 | 
 | 
| 5 |  | -This is still an experimental version and the following differences exist: | 
|  | 5 | +We simply expand the input channels of Flux.1 Dev from 64 to 128 to allow for additional inputs and then train a regular LoRA on top of it. To account for the newly added input channels, we additional append a LoRA on the underlying layer (`x_embedder`). Inference, however, is performed with the `FluxControlPipeline`.  | 
| 6 | 6 | 
 | 
| 7 |  | -* No use of bias on `lora_B`. | 
| 8 |  | -* Mo updates on the norm scales. | 
|  | 7 | +> [!NOTE] | 
|  | 8 | +> **Gated model** | 
|  | 9 | +> | 
|  | 10 | +> As the model is gated, before using it with diffusers you first need to go to the [FLUX.1 [dev] Hugging Face page](https://huggingface.co/black-forest-labs/FLUX.1-dev), fill in the form and accept the gate. Once you are in, you need to log in so that your system knows you’ve accepted the gate. Use the command below to log in: | 
| 9 | 11 | 
 | 
| 10 |  | -We simply expand the input channels of Flux.1 Dev from 64 to 128 to allow for additional inputs and then train a regular LoRA on top of it. To account for the newly added input channels, we additional append a LoRA on the underlying layer (`x_embedder`). Inference, however, is performed with the `FluxControlPipeline`.  | 
|  | 12 | +```bash | 
|  | 13 | +huggingface-cli login | 
|  | 14 | +``` | 
| 11 | 15 | 
 | 
| 12 | 16 | Example command: | 
| 13 | 17 | 
 | 
| @@ -35,4 +39,46 @@ accelerate launch train_control_lora_flux.py \ | 
| 35 | 39 | 
 | 
| 36 | 40 | `openpose.png` comes from [here](https://huggingface.co/Adapter/t2iadapter/resolve/main/openpose.png). | 
| 37 | 41 | 
 | 
| 38 |  | -You need to install `diffusers` from the branch of [this PR](https://github.com/huggingface/diffusers/pull/9999).  | 
|  | 42 | +You need to install `diffusers` from the branch of [this PR](https://github.com/huggingface/diffusers/pull/9999). When it's merged, you should install `diffusers` from the `main`. | 
|  | 43 | + | 
|  | 44 | +The training script exposes additional CLI args that might be useful to experiment with: | 
|  | 45 | + | 
|  | 46 | +* `use_lora_bias`: When set, additionally trains the biases of the `lora_B` layer.  | 
|  | 47 | +* `train_norm_layers`: When set, additionally trains the normalization scales. Takes care of saving and loading. | 
|  | 48 | +* `lora_layers`: Specify the layers you want to apply LoRA to. If you specify "all-linear", all the linear layers will be LoRA-attached. | 
|  | 49 | + | 
|  | 50 | +## Training with DeepSpeed | 
|  | 51 | + | 
|  | 52 | +It's possible to train with [DeepSpeed](https://github.com/microsoft/DeepSpeed), specifically leveraging the Zero2 system optimization. To use it, save the following config to an YAML file (feel free to modify as needed): | 
|  | 53 | + | 
|  | 54 | +```yaml | 
|  | 55 | +compute_environment: LOCAL_MACHINE | 
|  | 56 | +debug: false | 
|  | 57 | +deepspeed_config: | 
|  | 58 | +  gradient_accumulation_steps: 1 | 
|  | 59 | +  gradient_clipping: 1.0 | 
|  | 60 | +  offload_optimizer_device: cpu | 
|  | 61 | +  offload_param_device: cpu | 
|  | 62 | +  zero3_init_flag: false | 
|  | 63 | +  zero_stage: 2 | 
|  | 64 | +distributed_type: DEEPSPEED | 
|  | 65 | +downcast_bf16: 'no' | 
|  | 66 | +enable_cpu_affinity: false | 
|  | 67 | +machine_rank: 0 | 
|  | 68 | +main_training_function: main | 
|  | 69 | +mixed_precision: bf16 | 
|  | 70 | +num_machines: 1 | 
|  | 71 | +num_processes: 1 | 
|  | 72 | +rdzv_backend: static | 
|  | 73 | +same_network: true | 
|  | 74 | +tpu_env: [] | 
|  | 75 | +tpu_use_cluster: false | 
|  | 76 | +tpu_use_sudo: false | 
|  | 77 | +use_cpu: false | 
|  | 78 | +``` | 
|  | 79 | +
 | 
|  | 80 | +And then while launching training, pass the config file: | 
|  | 81 | +
 | 
|  | 82 | +```bash | 
|  | 83 | +accelerate launch --config_file=CONFIG_FILE.yaml ... | 
|  | 84 | +``` | 
0 commit comments