LoCon - LoRA for Convolution Network

Motivation

convolution is matmul so there should be a lora version for it.

Why Convolution is matmul?

im2col:

What I did

A demo for LoRA on Convolution network(This repo)
A network module for kohya_ss/sd-script
An Extension for using this↑ in sd-webui

Difference from training LoRA on Stable Diffusion

normally most of people train LoRA with kohya-ss/sd-scripts, (me too)

but lora only train green part, locon can train yellow part. Combine them can cover almost all of the layers in the model.

(I skip the porj in and porj out since they have very small channels, if you want to f/t them, maybe just f/t it without any trick.)

usage

For kohya script

Activate sd-scripts' venv and then install this package

source PATH_TO_SDSCRIPTS_VENV/Scripts/activate

PATH_TO_SDSCRIPTS_VENV\Scripts\Activate.ps1 # or .bat for cmd

And then you can install this package:

through pip

pip install locon

from source

git clone https://github.com/KohakuBlueleaf/LoCon
cd LoCon
pip install .

Finally you can use this package's locon_kohya to run kohya's training script

python3 sd-scripts/train_network.py \
  --network_module locon.locon_kohya \
  --network_dim "RANK_FOR_TRANSFORMER" --network_alpha "ALPHA_FOR_TRANSFORMER"\
  --network_args "conv_dim=RANK_FOR_CONV" "conv_alpha=ALPHA_FOR_CONV" "dropout=DROPOUT_RATE" \

to train locon+lora for SD model

For a1111's sd-webui

download Extension into sd-webui, and then use locon model as how you use lora model.

Additional Networks

Once you install the extension. You can also use your locon model in addnet
just use it as LoRA model.

Extract LoCon

You can extract LoCon from a dreambooth model with its base model.

python3 extract_locon.py <settings> <base_model> <db_model> <output>

Use --help to get more info

$ python3 extract_locon.py --help
usage: extract_locon.py [-h] [--is_v2] [--device DEVICE] [--mode MODE] [--safetensors] [--linear_dim LINEAR_DIM] [--conv_dim CONV_DIM]
                        [--linear_threshold LINEAR_THRESHOLD] [--conv_threshold CONV_THRESHOLD] [--linear_ratio LINEAR_RATIO] [--conv_ratio CONV_RATIO]
                        [--linear_percentile LINEAR_PERCENTILE] [--conv_percentile CONV_PERCENTILE]
                        base_model db_model output_name

Update

For main package

cd LoCon
git pull
pip install .

For extension
Use Extension's tab's check updates btn, and then restart the webui
Use apply and restart button will not restart the extension

Example Model

Yog-Sothoth LoRA/LoCon: https://civitai.com/models/14878/loconlora-yog-sothoth-depersonalization

LoRA rank=1: LoCon rank=1

xy grid:

Onimai LoRA: https://huggingface.co/KBlueLeaf/onimai-locon-test

Some calculation

LoRA for Linear:

$$Y_{out*batch} = W_{out*in}‧X_{in*batch}$$

$$Y'_{out*batch} = W_{out*in}‧X_{in*batch} + Wa_{out*rank}‧Wb_{rank*in}‧X_{in*batch}$$

Convolution img2col:

$$X:[channel, width, height]$$

$$\xrightarrow{reorder}[c*kw*kh, outw*outh]$$

$$Kernels: [out, c, kw, kh] \xrightarrow{reshape} [out, c*kw*kh]$$

$$Conv(X, Kernels) = Kernels * X \xrightarrow{reshape} [out, outw, outh]$$

LoRA for Convolution:

$$Conv(in, out, ksize, padding, stride)$$

$$\xrightarrow{}Conv(rank, out, 1)\circ Conv(in, rank, ksize, padding, stride)$$

Another form:

[out_ch, in_ch*size**2] x [in_ch*size**2, out_h * out_w]
↓
[out_ch, LoRA_rank] x [LoRA_rank, in_ch*size**2] x [in_ch*size**2, out_h * out_w]
↓
[out_ch, LoRA_rank] x ([LoRA_rank, in_ch*size**2] x [in_ch*size**2, out_h * out_w])
↓
[out_ch, LoRA_rank] x [LoRA_rank, out_h * out_w]

FLOPS:

before = out_ch * in_ch * size**2 * out_h * out_w
after = out_ch * LoRA_rank * out_h * out_w + LoRA_rank * in_ch * size**2 * out_h * out_w
after = (out_ch * LoRA_rank + LoRA_rank * in_ch * size**2) * out_h * out_w

Params to train:

before = out_ch * in_ch * size**2
after = LoRA_rank * in_ch * size**2 + LoRA_rank * out_ch

Citation

@misc{LoCon,
  author       = "Shih-Ying Yeh (Kohaku-BlueLeaf)",
  title        = "LoCon - LoRA for Convolution Network",
  howpublished = "\url{https://github.com/KohakuBlueleaf/LoCon}",
  month        = "Feb",
  year         = "2023"
}

253701/LoCon