Kandinsky-2 Documentation

To train Kandinsky 2.1 using the UnCLIP method, you must load the base configuration from the repository and override specific paths for the downloaded checkpoints, text encoder, and your training dataset.

Key configuration overrides:

• params_path: Path to the decoder checkpoint (e.g., kandinsky2/2_1/decoder_fp16.ckpt).
• image_enc_params['ckpt_path']: Path to the MOVQ checkpoint.
• text_enc_params['model_path']: Path to the text encoder directory.
• data['train']['tokenizer_name']: Path to the text encoder directory.
• data['train']['df_path']: Path to your CSV file containing image_name and caption columns.
• save_path: Directory where training checkpoints will be saved.

from omegaconf import OmegaConf

config = OmegaConf.load("Kandinsky-2/train_configs/config_unclip_2_1.yaml")
config['params_path'] = 'kandinsky2/2_1/decoder_fp16.ckpt'
config['image_enc_params']['ckpt_path'] = 'kandinsky2/2_1/movq_final.ckpt'
config['text_enc_params']['model_path'] = 'kandinsky2/2_1/text_encoder'
config['data']['train']['tokenizer_name'] = 'kandinsky2/2_1/text_encoder'
config['data']['train']['df_path'] = 'test.csv'
config['num_epochs'] = 1001
config['save_path'] = 'saves'

with open("train_config.yaml", "w") as f:
    OmegaConf.save(config, f)

Use the get_kandinsky2_1 function to download the necessary model files from Hugging Face. This includes the prior_fp16.ckpt weights and the ViT-L-14_stats.th statistics file required for the model to function. The files are stored in a subdirectory named 2_1 within your specified cache_dir.

def get_kandinsky2_1(
    cache_dir="./kandinsky2",
    use_auth_token=None,
):
    cache_dir = os.path.join(cache_dir, "2_1")

    prior_name = "prior_fp16.ckpt"
    config_file_url = hf_hub_url(repo_id="sberbank-ai/Kandinsky_2.1", filename=prior_name)
    cached_download(
        config_file_url,
        cache_dir=cache_dir,
        force_filename=prior_name,
        use_auth_token=use_auth_token,
    )

    config_file_url = hf_hub_url(repo_id="sberbank-ai/Kandinsky_2.1", filename="ViT-L-14_stats.th")
    cached_download(
        config_file_url,
        cache_dir=cache_dir,
        force_filename="ViT-L-14_stats.th",
        use_auth_token=use_auth_token,
    )

get_kandinsky2_1()

Kandinsky 2.2 requires two main components: a KandinskyV22PriorPipeline and a KandinskyV22Pipeline (decoder). You can optimize performance by loading the image_encoder and unet separately and passing them into the pipelines. It is recommended to use torch.float16 for efficiency on GPU.

import torch
from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
from transformers import CLIPVisionModelWithProjection
from diffusers.models import UNet2DConditionModel

DEVICE = torch.device('cuda') # Replace with your device

# 1. Load Image Encoder for the Prior
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
    'kandinsky-community/kandinsky-2-2-prior',
    subfolder='image_encoder'
).half().to(DEVICE)

# 2. Load UNet for the Decoder
unet = UNet2DConditionModel.from_pretrained(
    'kandinsky-community/kandinsky-2-2-decoder', 
    subfolder='unet'
).half().to(DEVICE)

# 3. Initialize Prior Pipeline
prior = KandinskyV22PriorPipeline.from_pretrained(
    'kandinsky-community/kandinsky-2-2-prior',
    image_encoder=image_encoder, 
    torch_dtype=torch.float16
).to(DEVICE)

# 4. Initialize Decoder Pipeline
decoder = KandinskyV22Pipeline.from_pretrained(
    'kandinsky-community/kandinsky-2-2-decoder',
    unet=unet, 
    torch_dtype=torch.float16
).to(DEVICE)

To train the prior, load the base configuration from Kandinsky-2/train_configs/config_prior.yaml and override the following keys:

• params_path: Path to the prior_fp16.ckpt file.
• clip_mean_std_path: Path to the ViT-L-14_stats.th file.
• data.train.df_path: Path to your training CSV.
• num_epochs: Number of training epochs.
• save_path: Directory where training checkpoints will be saved.

After saving the modified configuration to a file (e.g., train_config.yaml), run the training script using the --config flag.

from omegaconf import OmegaConf
import os

# 1. Load base config
config = OmegaConf.load("Kandinsky-2/train_configs/config_prior.yaml")

# 2. Override parameters
config['params_path'] = 'kandinsky2/2_1/prior_fp16.ckpt'
config['clip_mean_std_path'] = 'kandinsky2/2_1/ViT-L-14_stats.th'
config['data']['train']['df_path'] = 'test.csv'
config['num_epochs'] = 1001
config['save_path'] = 'saves'

if not os.path.exists('saves'):
    os.mkdir('saves')

# 3. Save the new config
with open("train_config.yaml", "w") as f:
    OmegaConf.save(config, f)

# 4. Run training via CLI
# !python Kandinsky-2/train_prior.py --config train_config.yaml