Kandinsky 2.2

Ce script est expérimental, et il est facile de sur-apprendre et de rencontrer des problèmes tels que l'oubli catastrophique. Essayez d'explorer différents hyperparamètres pour obtenir les meilleurs résultats sur votre jeu de données.

Kandinsky 2.2 est un modèle de génération d'images à partir de texte multilingue capable de produire des images plus photoréalistes. Le modèle comprend un modèle de prior d'image pour créer des embeddings d'image à partir de prompts textuels, et un modèle décodeur qui génère des images sur la base des embeddings du modèle de prior. C'est pourquoi vous trouverez deux scripts distincts dans Diffusers pour Kandinsky 2.2, l'un pour l'entraînement du modèle de prior et l'autre pour l'entraînement du modèle décodeur. Vous pouvez entraîner les deux modèles séparément, mais pour obtenir les meilleurs résultats, vous devez entraîner à la fois les modèles de prior et de décodeur.

Selon votre GPU, vous devrez peut-être activer gradient_checkpointing (⚠️ non pris en charge pour le modèle précédent !), mixed_precision et gradient_accumulation_steps pour aider à faire tenir le modèle en mémoire et accélérer l’entraînement. Vous pouvez réduire encore davantage votre consommation de mémoire en activant l’attention efficace en mémoire avec xFormers (la version v0.0.16 échoue pour l’entraînement sur certains GPU, il est donc possible que vous deviez installer une version de développement à la place).

Ce guide explore les scripts train_text_to_image_prior.py et train_text_to_image_decoder.py pour vous aider à mieux les connaître et à savoir comment les adapter à votre propre cas d'usage.

Avant d'exécuter les scripts, assurez-vous d'installer la bibliothèque à partir du code source :

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

Puis accédez au dossier d'exemple contenant le script d'entraînement et installez les dépendances requises pour le script que vous utilisez :

cd examples/kandinsky2_2/text_to_image
pip install -r requirements.txt

🤗 Accelerate est une bibliothèque qui vous aide à entraîner sur plusieurs GPU/TPU ou avec une précision mixte. Elle configure automatiquement votre configuration d’entraînement en fonction de votre matériel et de votre environnement. Consultez le tour rapide de 🤗 Accelerate pour en savoir plus.

Initialiser un environnement 🤗 Accelerate :

accelerate config

Pour configurer un environnement 🤗 Accelerate par défaut sans choisir de configuration :

accelerate config default

Sinon, si votre environnement ne prend pas en charge un shell interactif, comme un notebook, vous pouvez utiliser :

from accelerate.utils import write_basic_config

write_basic_config()

Enfin, si vous souhaitez entraîner un modèle sur votre propre jeu de données, consultez le guide Create a dataset for training pour apprendre à créer un jeu de données compatible avec le script d'entraînement.

Les sections suivantes mettent en évidence les parties des scripts d’entraînement importantes pour comprendre comment les modifier, mais elles ne couvrent pas tous les aspects des scripts en détail. Si vous souhaitez en savoir plus, n’hésitez pas à lire les scripts et à nous faire part de vos questions ou préoccupations.

Paramètres du script

Les scripts d’entraînement fournissent de nombreux paramètres pour vous aider à personnaliser votre session d’entraînement. Tous les paramètres et leurs descriptions se trouvent dans la fonction parse_args(). Les scripts d’entraînement fournissent des valeurs par défaut pour chaque paramètre, telles que la taille du lot d’entraînement et le taux d’apprentissage, mais vous pouvez également définir vos propres valeurs dans la commande d’entraînement si vous le souhaitez.

Par exemple, pour accélérer l’entraînement avec la précision mixte en utilisant le format fp16, ajoutez le paramètre --mixed_precision à la commande d’entraînement :

accelerate launch train_text_to_image_prior.py \
  --mixed_precision="fp16"

La plupart des paramètres sont identiques à ceux du guide d’entraînement Text-to-image, passons donc directement à la présentation des scripts d’entraînement Kandinsky !

Pondération Min-SNR

La stratégie de pondération Min-SNR peut aider à l'entraînement en rééquilibrant la perte pour atteindre une convergence plus rapide. Le script d'entraînement prend en charge la prédiction de epsilon (bruit) ou de v_prediction, mais Min-SNR est compatible avec les deux types de prédiction. Cette stratégie de pondération n'est prise en charge que par PyTorch.

Ajoutez le paramètre --snr_gamma et définissez-le à la valeur recommandée de 5,0 :

accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0

Script d'entraînement

Le script d’entraînement est également similaire au guide d’entraînement Text-to-image, mais il a été modifié pour prendre en charge l’entraînement des modèles prior et decoder. Ce guide se concentre sur le code propre aux scripts d’entraînement de Kandinsky 2.2.

modèle a priori

modèle décodeur

La fonction main() contient le code permettant de préparer le jeu de données et d’entraîner le modèle.

L'une des principales différences que vous remarquerez immédiatement est que le script d'entraînement charge également un CLIPImageProcessor - en plus d'un scheduler et d'un tokenizer - pour le prétraitement des images, ainsi qu'un modèle CLIPVisionModelWithProjection pour l'encodage des images :

noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
    args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", dtype=weight_dtype
    ).eval()

Kandinsky utilise un PriorTransformer pour générer les embeddings d’image, il est donc nécessaire de configurer l’optimiseur pour apprendre les paramètres du mode prior.

prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Ensuite, les légendes d'entrée sont tokenisées et les images sont prétraitées par le CLIPImageProcessor :

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
    return examples

Enfin, la boucle d’entraînement convertit les images d’entrée en latents, ajoute du bruit aux embeddings d’images et effectue une prédiction :

model_pred = prior(
    noisy_latents,
    timestep=timesteps,
    proj_embedding=prompt_embeds,
    encoder_hidden_states=text_encoder_hidden_states,
    attention_mask=text_mask,
).predicted_image_embedding

Si vous souhaitez en savoir plus sur le fonctionnement de la boucle d’entraînement, consultez le tutoriel Comprendre les pipelines, les modèles et les planificateurs, qui détaille le schéma de base du processus de débruitage.

Lancer le script

Une fois que vous avez apporté toutes vos modifications ou que vous êtes satisfait de la configuration par défaut, vous êtes prêt à lancer le script d’entraînement ! 🚀

Vous entraînerez sur le jeu de données Naruto BLIP captions pour générer vos propres personnages de Naruto, mais vous pouvez également créer et entraîner sur votre propre jeu de données en suivant le guide Create a dataset for training. Définissez la variable d’environnement DATASET_NAME avec le nom du jeu de données sur le Hub, ou, si vous entraînez sur vos propres fichiers, définissez la variable d’environnement TRAIN_DIR avec le chemin vers votre jeu de données.

Si vous entraînez sur plus d’un GPU, ajoutez le paramètre --multi_gpu à la commande accelerate launch.

Pour suivre la progression de l’entraînement avec Weights & Biases, ajoutez le paramètre --report_to=wandb à la commande d’entraînement. Vous devez également ajouter --validation_prompt à la commande d’entraînement afin de suivre les résultats. Cela peut s’avérer très utile pour déboguer le modèle et consulter les résultats intermédiaires.

modèle a priori

modèle décodeur

Une fois l’entraînement terminé, vous pouvez utiliser votre nouveau modèle pour l’inférence !

modèle a priori

modèle décodeur

Étapes suivantes

Félicitations pour l'entraînement d'un modèle Kandinsky 2.2 ! Pour en savoir plus sur l'utilisation de votre nouveau modèle, les guides suivants peuvent vous être utiles :

Consultez le guide Kandinsky pour apprendre à l’utiliser pour diverses tâches (texte-vers-image, image-vers-image, inpainting, interpolation) et à le combiner avec un ControlNet.
Consultez les guides d’entraînement DreamBooth et LoRA pour apprendre à entraîner un modèle Kandinsky personnalisé à partir de quelques images d’exemple seulement. Ces deux techniques d’entraînement peuvent même être combinées !

Update on GitHub