feat: add controlnet union model

This commit is contained in:
Dowon
2024-07-17 20:27:56 +09:00
parent ab845cf596
commit db3be81c4e
3 changed files with 9 additions and 4 deletions

View File

@@ -2,6 +2,7 @@ from __future__ import annotations
from dataclasses import dataclass
from functools import partial
from itertools import chain
from types import SimpleNamespace
from typing import Any
@@ -42,6 +43,9 @@ else:
"depth": ["depth_midas", "depth_hand_refiner"],
}
union = list(chain.from_iterable(cn_module_choices.values()))
cn_module_choices["union"] = union
class Widgets(SimpleNamespace):
def tolist(self):

View File

@@ -8,4 +8,5 @@ cn_model_module = {
"tile": "tile_resample",
"depth": "depth_midas",
}
cn_model_regex = re.compile("|".join(cn_model_module.keys()), flags=re.IGNORECASE)
_names = [*cn_model_module, "union"]
cn_model_regex = re.compile("|".join(_names), flags=re.IGNORECASE)

View File

@@ -61,13 +61,13 @@ class ControlNetExt:
if (not self.cn_available) or model == "None":
return
if module is None or module == "None":
if module == "None":
module = None
if module is None:
for m, v in cn_model_module.items():
if m in model:
module = v
break
else:
module = None
cn_units = [
self.external_cn.ControlNetUnit(