Merge pull request kohya-ss#1284 from sdbds/fix_traincontrolnet

kohya-ss · web-flow · commit de0e0b946815 · 2024-05-19T17:39:15.000+09:00
Fix train controlnet
diff --git a/requirements.txt b/requirements.txt
@@ -17,6 +17,7 @@ easygui==0.98.3
 toml==0.10.2
 voluptuous==0.13.1
 huggingface-hub==0.20.1
+omegaconf==2.3.0
 # for Image utils
 imagesize==1.4.1
 # for BLIP captioning
diff --git a/train_controlnet.py b/train_controlnet.py
@@ -5,7 +5,7 @@
 import random
 import time
 from multiprocessing import Value
-from types import SimpleNamespace
+from omegaconf import OmegaConf
 import toml
 
 from tqdm import tqdm
@@ -148,8 +148,10 @@ def train(args):
             "in_channels": 4,
             "layers_per_block": 2,
             "mid_block_scale_factor": 1,
+            "mid_block_type": "UNetMidBlock2DCrossAttn",
             "norm_eps": 1e-05,
             "norm_num_groups": 32,
+            "num_attention_heads": [5, 10, 20, 20],
             "num_class_embeds": None,
             "only_cross_attention": False,
             "out_channels": 4,
@@ -179,8 +181,10 @@ def train(args):
             "in_channels": 4,
             "layers_per_block": 2,
             "mid_block_scale_factor": 1,
+            "mid_block_type": "UNetMidBlock2DCrossAttn",
             "norm_eps": 1e-05,
             "norm_num_groups": 32,
+            "num_attention_heads": 8,
             "out_channels": 4,
             "sample_size": 64,
             "up_block_types": ["UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"],
@@ -193,7 +197,7 @@ def train(args):
             "resnet_time_scale_shift": "default",
             "projection_class_embeddings_input_dim": None,
         }
-    unet.config = SimpleNamespace(**unet.config)
+    unet.config = OmegaConf.create(unet.config)
 
     controlnet = ControlNetModel.from_unet(unet)
 
