File size: 2,595 Bytes
7bed60d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
from __future__ import annotations
import copy
import numpy as np
from lib_controlnet import external_code, global_state
from lib_controlnet.external_code import ControlNetUnit
from modules import scripts
from modules.processing import StableDiffusionProcessing
from .common import cn_model_regex
controlnet_exists = True
controlnet_type = "forge"
def find_script(p: StableDiffusionProcessing, script_title: str) -> scripts.Script:
script = next((s for s in p.scripts.scripts if s.title() == script_title), None)
if not script:
msg = f"Script not found: {script_title!r}"
raise RuntimeError(msg)
return script
def add_forge_script_to_adetailer_run(
p: StableDiffusionProcessing, script_title: str, script_args: list
):
p.scripts = copy.copy(scripts.scripts_img2img)
p.scripts.alwayson_scripts = []
p.script_args_value = []
script = copy.copy(find_script(p, script_title))
script.args_from = len(p.script_args_value)
script.args_to = len(p.script_args_value) + len(script_args)
p.scripts.alwayson_scripts.append(script)
p.script_args_value.extend(script_args)
class ControlNetExt:
def __init__(self):
self.cn_available = False
self.external_cn = external_code
def init_controlnet(self):
self.cn_available = True
def update_scripts_args(
self,
p,
model: str,
module: str | None,
weight: float,
guidance_start: float,
guidance_end: float,
):
if (not self.cn_available) or model == "None":
return
image = np.asarray(p.init_images[0])
mask = np.full_like(image, fill_value=255)
cnet_image = {"image": image, "mask": mask}
pres = external_code.pixel_perfect_resolution(
image,
target_H=p.height,
target_W=p.width,
resize_mode=external_code.resize_mode_from_value(p.resize_mode),
)
add_forge_script_to_adetailer_run(
p,
"ControlNet",
[
ControlNetUnit(
enabled=True,
image=cnet_image,
model=model,
module=module,
weight=weight,
guidance_start=guidance_start,
guidance_end=guidance_end,
processor_res=pres,
)
],
)
def get_cn_models() -> list[str]:
models = global_state.get_all_controlnet_names()
return [m for m in models if cn_model_regex.search(m)]
|