1 files changed, 14 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index d7878cd..082e9b7 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -13,10 +13,12 @@ from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
 from slugify import slugify
+from timm.models import create_model
 import transformers
 
 from util.files import load_config, load_embeddings_from_dir
 from data.csv import VlpnDataModule, keyword_filter
+from models.convnext.discriminator import ConvNeXtDiscriminator
 from training.functional import train, add_placeholder_tokens, get_models
 from training.strategy.ti import textual_inversion_strategy
 from training.optimization import get_scheduler
@@ -661,6 +663,17 @@ def main():
         unet.enable_gradient_checkpointing()
         text_encoder.gradient_checkpointing_enable()
 
+    convnext = create_model(
+        "convnext_tiny",
+        pretrained=False,
+        num_classes=3,
+        drop_path_rate=0.0,
+    )
+    convnext.to(accelerator.device, dtype=weight_dtype)
+    convnext.requires_grad_(False)
+    convnext.eval()
+    disc = ConvNeXtDiscriminator(convnext, input_size=384)
+
     if len(args.alias_tokens) != 0:
         alias_placeholder_tokens = args.alias_tokens[::2]
         alias_initializer_tokens = args.alias_tokens[1::2]
@@ -802,6 +815,7 @@ def main():
         milestone_checkpoints=not args.no_milestone_checkpoints,
         global_step_offset=global_step_offset,
         offset_noise_strength=args.offset_noise_strength,
+        disc=disc,
         # --
         use_emb_decay=args.use_emb_decay,
         emb_decay_target=args.emb_decay_target,