mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
576 lines
29 KiB
Python
576 lines
29 KiB
Python
import os
|
|
import time
|
|
import threading
|
|
from typing import Union
|
|
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, FluxPipeline, StableDiffusion3Pipeline, ControlNetModel
|
|
from modules.control.units import detect
|
|
from modules.shared import log, opts, cmd_opts, state, listdir
|
|
from modules import errors, sd_models, devices, model_quant
|
|
from modules.processing import StableDiffusionProcessingControl
|
|
|
|
|
|
what = 'ControlNet'
|
|
debug = os.environ.get('SD_CONTROL_DEBUG', None) is not None
|
|
debug_log = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
|
|
predefined_sd15 = {
|
|
'Canny': "lllyasviel/control_v11p_sd15_canny",
|
|
'Depth': "lllyasviel/control_v11f1p_sd15_depth",
|
|
'HED': "lllyasviel/sd-controlnet-hed",
|
|
'IP2P': "lllyasviel/control_v11e_sd15_ip2p",
|
|
'LineArt': "lllyasviel/control_v11p_sd15_lineart",
|
|
'LineArt Anime': "lllyasviel/control_v11p_sd15s2_lineart_anime",
|
|
'MLDS': "lllyasviel/control_v11p_sd15_mlsd",
|
|
'NormalBae': "lllyasviel/control_v11p_sd15_normalbae",
|
|
'OpenPose': "lllyasviel/control_v11p_sd15_openpose",
|
|
'Scribble': "lllyasviel/control_v11p_sd15_scribble",
|
|
'Segment': "lllyasviel/control_v11p_sd15_seg",
|
|
'Shuffle': "lllyasviel/control_v11e_sd15_shuffle",
|
|
'SoftEdge': "lllyasviel/control_v11p_sd15_softedge",
|
|
'Tile': "lllyasviel/control_v11f1e_sd15_tile",
|
|
'Depth Anything': 'vladmandic/depth-anything',
|
|
'Canny FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_canny.safetensors',
|
|
'Inpaint FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_inpaint.safetensors',
|
|
'LineArt Anime FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_animeline.safetensors',
|
|
'LineArt FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_lineart.safetensors',
|
|
'MLSD FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_mlsd.safetensors',
|
|
'NormalBae FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_normal.safetensors',
|
|
'OpenPose FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_openpose.safetensors',
|
|
'Pix2Pix FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_pix2pix.safetensors',
|
|
'Scribble FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_scribble.safetensors',
|
|
'Segment FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_seg.safetensors',
|
|
'Shuffle FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_shuffle.safetensors',
|
|
'SoftEdge FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_softedge.safetensors',
|
|
'Tile FP16': 'Aptronym/SDNext/ControlNet11/controlnet11Models_tileE.safetensors',
|
|
'CiaraRowles TemporalNet': "CiaraRowles/TemporalNet",
|
|
'Ciaochaos Recolor': 'ioclab/control_v1p_sd15_brightness',
|
|
'Ciaochaos Illumination': 'ioclab/control_v1u_sd15_illumination/illumination20000.safetensors',
|
|
}
|
|
predefined_sdxl = {
|
|
'Canny Small XL': 'diffusers/controlnet-canny-sdxl-1.0-small',
|
|
'Canny Mid XL': 'diffusers/controlnet-canny-sdxl-1.0-mid',
|
|
'Canny XL': 'diffusers/controlnet-canny-sdxl-1.0',
|
|
'Depth Zoe XL': 'diffusers/controlnet-zoe-depth-sdxl-1.0',
|
|
'Depth Mid XL': 'diffusers/controlnet-depth-sdxl-1.0-mid',
|
|
'OpenPose XL': 'thibaud/controlnet-openpose-sdxl-1.0/bin',
|
|
'Xinsir Union XL': 'xinsir/controlnet-union-sdxl-1.0',
|
|
'Xinsir ProMax XL': 'brad-twinkl/controlnet-union-sdxl-1.0-promax',
|
|
'Xinsir OpenPose XL': 'xinsir/controlnet-openpose-sdxl-1.0',
|
|
'Xinsir Canny XL': 'xinsir/controlnet-canny-sdxl-1.0',
|
|
'Xinsir Depth XL': 'xinsir/controlnet-depth-sdxl-1.0',
|
|
'Xinsir Scribble XL': 'xinsir/controlnet-scribble-sdxl-1.0',
|
|
'Xinsir Anime Painter XL': 'xinsir/anime-painter',
|
|
'Xinsir Tile XL': 'xinsir/controlnet-tile-sdxl-1.0',
|
|
'NoobAI Canny XL': 'Eugeoter/noob-sdxl-controlnet-canny',
|
|
'NoobAI Lineart Anime XL': 'Eugeoter/noob-sdxl-controlnet-lineart_anime',
|
|
'NoobAI Depth XL': 'Eugeoter/noob-sdxl-controlnet-depth',
|
|
'NoobAI Normal XL': 'Eugeoter/noob-sdxl-controlnet-normal',
|
|
'NoobAI SoftEdge XL': 'Eugeoter/noob-sdxl-controlnet-softedge_hed',
|
|
'NoobAI OpenPose XL': 'einar77/noob-openpose',
|
|
'TTPlanet Tile Realistic XL': 'Yakonrus/SDXL_Controlnet_Tile_Realistic_v2',
|
|
# 'StabilityAI Canny R128': 'stabilityai/control-lora/control-LoRAs-rank128/control-lora-canny-rank128.safetensors',
|
|
# 'StabilityAI Depth R128': 'stabilityai/control-lora/control-LoRAs-rank128/control-lora-depth-rank128.safetensors',
|
|
# 'StabilityAI Recolor R128': 'stabilityai/control-lora/control-LoRAs-rank128/control-lora-recolor-rank128.safetensors',
|
|
# 'StabilityAI Sketch R128': 'stabilityai/control-lora/control-LoRAs-rank128/control-lora-sketch-rank128-metadata.safetensors',
|
|
# 'StabilityAI Canny R256': 'stabilityai/control-lora/control-LoRAs-rank256/control-lora-canny-rank256.safetensors',
|
|
# 'StabilityAI Depth R256': 'stabilityai/control-lora/control-LoRAs-rank256/control-lora-depth-rank256.safetensors',
|
|
# 'StabilityAI Recolor R256': 'stabilityai/control-lora/control-LoRAs-rank256/control-lora-recolor-rank256.safetensors',
|
|
# 'StabilityAI Sketch R256': 'stabilityai/control-lora/control-LoRAs-rank256/control-lora-sketch-rank256.safetensors',
|
|
}
|
|
predefined_f1 = {
|
|
"InstantX Union F1": 'InstantX/FLUX.1-dev-Controlnet-Union',
|
|
"InstantX Canny F1": 'InstantX/FLUX.1-dev-Controlnet-Canny',
|
|
"JasperAI Depth F1": 'jasperai/Flux.1-dev-Controlnet-Depth',
|
|
"BlackForrestLabs Canny LoRA F1": '/huggingface.co/black-forest-labs/FLUX.1-Canny-dev-lora/flux1-canny-dev-lora.safetensors',
|
|
"BlackForrestLabs Depth LoRA F1": '/huggingface.co/black-forest-labs/FLUX.1-Depth-dev-lora/flux1-depth-dev-lora.safetensors',
|
|
"JasperAI Surface Normals F1": 'jasperai/Flux.1-dev-Controlnet-Surface-Normals',
|
|
"JasperAI Upscaler F1": 'jasperai/Flux.1-dev-Controlnet-Upscaler',
|
|
"Shakker-Labs Union F1": 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro',
|
|
"Shakker-Labs Pose F1": 'Shakker-Labs/FLUX.1-dev-ControlNet-Pose',
|
|
"Shakker-Labs Depth F1": 'Shakker-Labs/FLUX.1-dev-ControlNet-Depth',
|
|
"XLabs-AI Canny F1": 'XLabs-AI/flux-controlnet-canny-diffusers',
|
|
"XLabs-AI Depth F1": 'XLabs-AI/flux-controlnet-depth-diffusers',
|
|
"XLabs-AI HED F1": 'XLabs-AI/flux-controlnet-hed-diffusers',
|
|
"LibreFlux Segment F1": 'neuralvfx/LibreFlux-ControlNet',
|
|
}
|
|
predefined_sd3 = {
|
|
"StabilityAI Canny SD35": 'diffusers-internal-dev/sd35-controlnet-canny-8b',
|
|
"StabilityAI Depth SD35": 'diffusers-internal-dev/sd35-controlnet-depth-8b',
|
|
"StabilityAI Blur SD35": 'diffusers-internal-dev/sd35-controlnet-blur-8b',
|
|
"InstantX Canny SD35": 'InstantX/SD3-Controlnet-Canny',
|
|
"InstantX Pose SD35": 'InstantX/SD3-Controlnet-Pose',
|
|
"InstantX Depth SD35": 'InstantX/SD3-Controlnet-Depth',
|
|
"InstantX Tile SD35": 'InstantX/SD3-Controlnet-Tile',
|
|
"Alimama Inpainting SD35": 'alimama-creative/SD3-Controlnet-Inpainting',
|
|
"Alimama SoftEdge SD35": 'alimama-creative/SD3-Controlnet-Softedge',
|
|
}
|
|
predefined_qwen = {
|
|
"InstantX Union Qwen": 'InstantX/Qwen-Image-ControlNet-Union',
|
|
}
|
|
predefined_hunyuandit = {
|
|
"HunyuanDiT Canny": 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Canny',
|
|
"HunyuanDiT Pose": 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Pose',
|
|
"HunyuanDiT Depth": 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Depth',
|
|
}
|
|
predefined_zimage = {
|
|
"Z-Image-Turbo Union 1.0": 'hlky/Z-Image-Turbo-Fun-Controlnet-Union',
|
|
"Z-Image-Turbo Union 2.0": 'hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0',
|
|
"Z-Image-Turbo Union 2.1": 'hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1',
|
|
}
|
|
|
|
variants = {
|
|
'NoobAI Canny XL': 'fp16',
|
|
'NoobAI Lineart Anime XL': 'fp16',
|
|
'NoobAI Depth XL': 'fp16',
|
|
'NoobAI Normal XL': 'fp16',
|
|
'NoobAI SoftEdge XL': 'fp16',
|
|
'TTPlanet Tile Realistic XL': 'fp16',
|
|
}
|
|
|
|
subfolders = {
|
|
"LibreFlux Segment F1": 'controlnet',
|
|
}
|
|
|
|
remote_code = {
|
|
"LibreFlux Segment F1": True,
|
|
}
|
|
|
|
models = {}
|
|
all_models = {}
|
|
all_models.update(predefined_sd15)
|
|
all_models.update(predefined_sdxl)
|
|
all_models.update(predefined_f1)
|
|
all_models.update(predefined_sd3)
|
|
all_models.update(predefined_qwen)
|
|
all_models.update(predefined_hunyuandit)
|
|
all_models.update(predefined_zimage)
|
|
cache_dir = 'models/control/controlnet'
|
|
load_lock = threading.Lock()
|
|
|
|
|
|
def find_models():
|
|
path = os.path.join(opts.control_dir, 'controlnet')
|
|
files = listdir(path)
|
|
folders = [f for f in files if os.path.isdir(f) if os.path.exists(os.path.join(f, 'config.json'))]
|
|
files = [f for f in files if f.endswith('.safetensors')]
|
|
downloaded_models = {}
|
|
for f in files:
|
|
basename = os.path.splitext(os.path.relpath(f, path))[0]
|
|
downloaded_models[basename] = f
|
|
for f in folders:
|
|
basename = os.path.relpath(f, path)
|
|
downloaded_models[basename] = f
|
|
all_models.update(downloaded_models)
|
|
return downloaded_models
|
|
|
|
find_models()
|
|
|
|
|
|
def api_list_models(model_type: str = None):
|
|
import modules.shared
|
|
model_type = model_type or modules.shared.sd_model_type
|
|
model_list = []
|
|
if model_type == 'sd' or model_type == 'all':
|
|
model_list += list(predefined_sd15)
|
|
if model_type == 'sdxl' or model_type == 'all':
|
|
model_list += list(predefined_sdxl)
|
|
if model_type == 'f1' or model_type == 'all':
|
|
model_list += list(predefined_f1)
|
|
if model_type == 'sd3' or model_type == 'all':
|
|
model_list += list(predefined_sd3)
|
|
if model_type == 'qwen' or model_type == 'all':
|
|
model_list += list(predefined_qwen)
|
|
if model_type == 'hunyuandit' or model_type == 'all':
|
|
model_list += list(predefined_hunyuandit)
|
|
if model_type == 'zimage':
|
|
model_list += list(predefined_zimage)
|
|
model_list += sorted(find_models())
|
|
return model_list
|
|
|
|
|
|
def list_models(refresh=False):
|
|
import modules.shared
|
|
global models # pylint: disable=global-statement
|
|
if not refresh and len(models) > 0:
|
|
return models
|
|
models = {}
|
|
if modules.shared.sd_model_type == 'none':
|
|
models = ['None']
|
|
elif modules.shared.sd_model_type == 'sdxl':
|
|
models = ['None'] + list(predefined_sdxl) + sorted(find_models())
|
|
elif modules.shared.sd_model_type == 'sd':
|
|
models = ['None'] + list(predefined_sd15) + sorted(find_models())
|
|
elif modules.shared.sd_model_type == 'f1':
|
|
models = ['None'] + list(predefined_f1) + sorted(find_models())
|
|
elif modules.shared.sd_model_type == 'sd3':
|
|
models = ['None'] + list(predefined_sd3) + sorted(find_models())
|
|
elif modules.shared.sd_model_type == 'qwen':
|
|
models = ['None'] + list(predefined_qwen) + sorted(find_models())
|
|
elif modules.shared.sd_model_type == 'hunyuandit':
|
|
models = ['None'] + list(predefined_hunyuandit) + sorted(find_models())
|
|
elif modules.shared.sd_model_type == 'zimage':
|
|
models = ['None'] + list(predefined_zimage) + sorted(find_models())
|
|
else:
|
|
log.warning(f'Control {what} model list failed: unknown model type')
|
|
models = ['None'] + list(all_models) + sorted(find_models())
|
|
debug_log(f'Control list {what}: path={cache_dir} models={models}')
|
|
return models
|
|
|
|
|
|
class ControlNet():
|
|
def __init__(self, model_id: str = None, device = None, dtype = None, load_config = None):
|
|
self.model: ControlNetModel = None
|
|
self.model_id: str = model_id
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.load_config = { 'cache_dir': cache_dir }
|
|
if load_config is not None:
|
|
self.load_config.update(load_config)
|
|
if opts.offline_mode:
|
|
self.load_config["local_files_only"] = True
|
|
os.environ['HF_HUB_OFFLINE'] = '1'
|
|
else:
|
|
os.environ.pop('HF_HUB_OFFLINE', None)
|
|
os.unsetenv('HF_HUB_OFFLINE')
|
|
if model_id is not None:
|
|
self.load()
|
|
|
|
def __str__(self):
|
|
return f' ControlNet(id={self.model_id} model={self.model.__class__.__name__})' if self.model_id and self.model else ''
|
|
|
|
def reset(self):
|
|
if self.model is not None:
|
|
debug_log(f'Control {what} model unloaded')
|
|
self.model = None
|
|
self.model_id = None
|
|
devices.torch_gc(force=True, reason='controlnet')
|
|
|
|
def get_class(self, model_id:str=''):
|
|
from modules import shared
|
|
if shared.sd_model_type == 'none':
|
|
_load = shared.sd_model # trigger a load
|
|
if shared.sd_model_type == 'sd':
|
|
from diffusers import ControlNetModel as cls # pylint: disable=reimported
|
|
config = 'lllyasviel/control_v11p_sd15_canny'
|
|
elif shared.sd_model_type == 'sdxl':
|
|
if 'union' in model_id.lower():
|
|
from diffusers import ControlNetUnionModel as cls
|
|
config = 'xinsir/controlnet-union-sdxl-1.0'
|
|
elif 'promax' in model_id.lower():
|
|
from diffusers import ControlNetUnionModel as cls
|
|
config = 'brad-twinkl/controlnet-union-sdxl-1.0-promax'
|
|
else:
|
|
from diffusers import ControlNetModel as cls # pylint: disable=reimported # sdxl shares same model class
|
|
config = 'Eugeoter/noob-sdxl-controlnet-canny'
|
|
elif shared.sd_model_type == 'f1':
|
|
from diffusers import FluxControlNetModel as cls
|
|
config = 'InstantX/FLUX.1-dev-Controlnet-Union'
|
|
elif shared.sd_model_type == 'sd3':
|
|
from diffusers import SD3ControlNetModel as cls
|
|
config = 'InstantX/SD3-Controlnet-Canny'
|
|
elif shared.sd_model_type == 'qwen':
|
|
from diffusers import QwenImageControlNetModel as cls
|
|
config = 'InstantX/Qwen-Image-ControlNet-Union'
|
|
elif shared.sd_model_type == 'hunyuandit':
|
|
from diffusers import HunyuanDiT2DControlNetModel as cls
|
|
config = 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Canny'
|
|
elif shared.sd_model_type == 'zimage':
|
|
from diffusers import ZImageControlNetModel as cls
|
|
if '2.0' in model_id:
|
|
config = 'hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0'
|
|
elif '2.1' in model_id:
|
|
config = 'hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1'
|
|
else:
|
|
config = 'hlky/Z-Image-Turbo-Fun-Controlnet-Union'
|
|
else:
|
|
log.error(f'Control {what}: type={shared.sd_model_type} unsupported model')
|
|
return None, None
|
|
return cls, config
|
|
|
|
def load_safetensors(self, model_id, model_path, cls, config): # pylint: disable=unused-argument
|
|
name = os.path.splitext(model_path)[0]
|
|
config_path = None
|
|
if not os.path.exists(model_path):
|
|
import huggingface_hub as hf
|
|
parts = model_path.split('/')
|
|
repo_id = f'{parts[0]}/{parts[1]}'
|
|
filename = os.path.splitext('/'.join(parts[2:]))[0]
|
|
model_path = hf.hf_hub_download(repo_id=repo_id, filename=f'{filename}.safetensors', cache_dir=cache_dir)
|
|
if config_path is None:
|
|
try:
|
|
config_path = hf.hf_hub_download(repo_id=repo_id, filename=f'{filename}.yaml', cache_dir=cache_dir)
|
|
except Exception:
|
|
pass # no yaml file
|
|
if config_path is None:
|
|
try:
|
|
config_path = hf.hf_hub_download(repo_id=repo_id, filename=f'{filename}.json', cache_dir=cache_dir)
|
|
except Exception:
|
|
pass # no yaml file
|
|
elif os.path.exists(name + '.yaml'):
|
|
config_path = f'{name}.yaml'
|
|
elif os.path.exists(name + '.json'):
|
|
config_path = f'{name}.json'
|
|
if config_path is not None:
|
|
self.load_config['original_config_file '] = config_path
|
|
self.model = cls.from_single_file(model_path, config=config, **self.load_config)
|
|
|
|
def load(self, model_id: str = None, force: bool = False) -> str:
|
|
with load_lock:
|
|
try:
|
|
t0 = time.time()
|
|
model_id = model_id or self.model_id
|
|
if model_id is None or model_id == 'None':
|
|
self.reset()
|
|
return
|
|
if model_id not in all_models:
|
|
log.error(f'Control {what}: id="{model_id}" available={list(all_models)} unknown model')
|
|
return
|
|
model_path = all_models[model_id]
|
|
if model_path == '':
|
|
return
|
|
if model_path is None:
|
|
log.error(f'Control {what} model load: id="{model_id}" unknown model id')
|
|
return
|
|
if 'lora' in model_id.lower():
|
|
self.model = model_path
|
|
return
|
|
if model_id == self.model_id and not force:
|
|
# log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
|
|
return
|
|
log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}"')
|
|
cls, config = self.get_class(model_id)
|
|
if cls is None:
|
|
log.error(f'Control {what} model load: id="{model_id}" unknown base model')
|
|
return
|
|
self.reset()
|
|
jobid = state.begin(f'Load {what}')
|
|
if model_path.endswith('.safetensors'):
|
|
self.load_safetensors(model_id, model_path, cls, config)
|
|
else:
|
|
kwargs = {}
|
|
if '/bin' in model_path:
|
|
model_path = model_path.replace('/bin', '')
|
|
self.load_config['use_safetensors'] = False
|
|
else:
|
|
self.load_config['use_safetensors'] = True
|
|
if variants.get(model_id, None) is not None:
|
|
kwargs['variant'] = variants[model_id]
|
|
if subfolders.get(model_id, None) is not None:
|
|
kwargs['subfolder'] = subfolders[model_id]
|
|
if remote_code.get(model_id, None) is not None:
|
|
kwargs['trust_remote_code'] = remote_code[model_id]
|
|
try:
|
|
self.model = cls.from_pretrained(model_path, **self.load_config, **kwargs)
|
|
except Exception as e:
|
|
log.error(f'Control {what} model load: id="{model_id}" {e}')
|
|
if debug:
|
|
errors.display(e, 'Control')
|
|
if self.model is None:
|
|
return
|
|
if not cmd_opts.lowvram: # lowvram will cause unet<->controlnet to ping-pong but saves more memory
|
|
self.model.offload_never = True
|
|
if self.dtype is not None:
|
|
self.model.to(self.dtype)
|
|
if self.device is not None:
|
|
if (opts.diffusers_offload_mode != 'balanced') and hasattr(self.model, 'to'):
|
|
try:
|
|
self.model.to(self.device)
|
|
except Exception as e:
|
|
if 'Cannot copy out of meta tensor' in str(e):
|
|
self.model.to_empty(device=self.device)
|
|
if "Control" in opts.sdnq_quantize_weights:
|
|
try:
|
|
log.debug(f'Control {what} model SDNQ quantize: id="{model_id}"')
|
|
from modules.model_quant import sdnq_quantize_model
|
|
self.model = sdnq_quantize_model(self.model)
|
|
except Exception as e:
|
|
log.error(f'Control {what} model SDNQ Compression failed: id="{model_id}" {e}')
|
|
elif "Control" in opts.optimum_quanto_weights:
|
|
try:
|
|
log.debug(f'Control {what} model Optimum Quanto: id="{model_id}"')
|
|
model_quant.load_quanto('Load model: type=Control')
|
|
from modules.model_quant import optimum_quanto_model
|
|
self.model = optimum_quanto_model(self.model)
|
|
except Exception as e:
|
|
log.error(f'Control {what} model Optimum Quanto: id="{model_id}" {e}')
|
|
elif "Control" in opts.torchao_quantization:
|
|
try:
|
|
log.debug(f'Control {what} model Torch AO: id="{model_id}"')
|
|
model_quant.load_torchao('Load model: type=Control')
|
|
from modules.model_quant import torchao_quantization
|
|
self.model = torchao_quantization(self.model)
|
|
except Exception as e:
|
|
log.error(f'Control {what} model Torch AO: id="{model_id}" {e}')
|
|
if self.device is not None:
|
|
sd_models.move_model(self.model, self.device)
|
|
if "Control" in opts.cuda_compile:
|
|
try:
|
|
from modules.sd_models_compile import compile_torch
|
|
self.model = compile_torch(self.model, apply_to_components=False, op="Control")
|
|
except Exception as e:
|
|
log.warning(f"Control compile error: {e}")
|
|
t1 = time.time()
|
|
self.model_id = model_id
|
|
log.info(f'Control {what} model loaded: id="{self.model_id}" path="{model_path}" cls={cls.__name__} time={t1-t0:.2f}')
|
|
state.end(jobid)
|
|
return f'{what} loaded model: {self.model_id}'
|
|
except Exception as e:
|
|
log.error(f'Control {what} model load: id="{model_id}" {e}')
|
|
errors.display(e, f'Control {what} load')
|
|
return f'{what} failed to load model: {model_id}'
|
|
|
|
|
|
class ControlNetPipeline():
|
|
def __init__(self,
|
|
controlnet: Union[ControlNetModel, list[ControlNetModel]],
|
|
pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline, FluxPipeline, StableDiffusion3Pipeline],
|
|
dtype = None,
|
|
p: StableDiffusionProcessingControl = None, # pylint: disable=unused-argument
|
|
):
|
|
t0 = time.time()
|
|
self.orig_pipeline = pipeline
|
|
self.pipeline = None
|
|
|
|
controlnets = controlnet if isinstance(controlnet, list) else [controlnet]
|
|
loras = [cn for cn in controlnets if isinstance(cn, str)]
|
|
controlnets = [cn for cn in controlnets if not isinstance(cn, str)]
|
|
|
|
if pipeline is None:
|
|
log.error('Control model pipeline: model not loaded')
|
|
return
|
|
elif detect.is_sdxl(pipeline) and len(controlnets) > 0:
|
|
from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetUnionPipeline
|
|
classes = [c.__class__.__name__ for c in controlnets]
|
|
if any(c == 'ControlNetUnionModel' for c in classes):
|
|
if not all(c == 'ControlNetUnionModel' for c in classes):
|
|
log.warning(f'Control {what}: units={classes} mixed type is not supported')
|
|
return
|
|
if isinstance(controlnets, list) and len(controlnets) == 1:
|
|
controlnets = controlnets[0]
|
|
cls = StableDiffusionXLControlNetUnionPipeline
|
|
else:
|
|
cls = StableDiffusionXLControlNetPipeline
|
|
self.pipeline = cls(
|
|
vae=pipeline.vae,
|
|
text_encoder=pipeline.text_encoder,
|
|
text_encoder_2=pipeline.text_encoder_2,
|
|
tokenizer=pipeline.tokenizer,
|
|
tokenizer_2=pipeline.tokenizer_2,
|
|
unet=pipeline.unet,
|
|
scheduler=pipeline.scheduler,
|
|
feature_extractor=getattr(pipeline, 'feature_extractor', None),
|
|
image_encoder=getattr(pipeline, 'image_encoder', None),
|
|
controlnet=controlnets, # can be a list
|
|
)
|
|
elif detect.is_f1(pipeline) and len(controlnets) > 0:
|
|
from diffusers import FluxControlNetPipeline
|
|
self.pipeline = FluxControlNetPipeline(
|
|
vae=pipeline.vae.to(devices.device),
|
|
text_encoder=pipeline.text_encoder,
|
|
text_encoder_2=pipeline.text_encoder_2,
|
|
tokenizer=pipeline.tokenizer,
|
|
tokenizer_2=pipeline.tokenizer_2,
|
|
transformer=pipeline.transformer,
|
|
scheduler=pipeline.scheduler,
|
|
controlnet=controlnets, # can be a list
|
|
)
|
|
elif detect.is_sd3(pipeline) and len(controlnets) > 0:
|
|
from diffusers import StableDiffusion3ControlNetPipeline
|
|
self.pipeline = StableDiffusion3ControlNetPipeline(
|
|
vae=pipeline.vae,
|
|
text_encoder=pipeline.text_encoder,
|
|
text_encoder_2=pipeline.text_encoder_2,
|
|
text_encoder_3=pipeline.text_encoder_3,
|
|
tokenizer=pipeline.tokenizer,
|
|
tokenizer_2=pipeline.tokenizer_2,
|
|
tokenizer_3=pipeline.tokenizer_3,
|
|
transformer=pipeline.transformer,
|
|
scheduler=pipeline.scheduler,
|
|
controlnet=controlnets, # can be a list
|
|
)
|
|
elif detect.is_sd15(pipeline) and len(controlnets) > 0:
|
|
from diffusers import StableDiffusionControlNetPipeline
|
|
self.pipeline = StableDiffusionControlNetPipeline(
|
|
vae=pipeline.vae,
|
|
text_encoder=pipeline.text_encoder,
|
|
tokenizer=pipeline.tokenizer,
|
|
unet=pipeline.unet,
|
|
scheduler=pipeline.scheduler,
|
|
feature_extractor=getattr(pipeline, 'feature_extractor', None),
|
|
image_encoder=getattr(pipeline, 'image_encoder', None),
|
|
requires_safety_checker=False,
|
|
safety_checker=None,
|
|
controlnet=controlnets, # can be a list
|
|
)
|
|
sd_models.move_model(self.pipeline, pipeline.device)
|
|
elif detect.is_qwen(pipeline) and len(controlnets) > 0:
|
|
from diffusers import QwenImageControlNetPipeline
|
|
self.pipeline = QwenImageControlNetPipeline(
|
|
vae=pipeline.vae,
|
|
text_encoder=pipeline.text_encoder,
|
|
tokenizer=pipeline.tokenizer,
|
|
transformer=pipeline.transformer,
|
|
scheduler=pipeline.scheduler,
|
|
controlnet=controlnets[0] if isinstance(controlnets, list) else controlnets, # can be a list
|
|
)
|
|
elif detect.is_hunyuandit(pipeline) and len(controlnets) > 0:
|
|
from diffusers import HunyuanDiTControlNetPipeline
|
|
self.pipeline = HunyuanDiTControlNetPipeline(
|
|
vae=pipeline.vae,
|
|
text_encoder=pipeline.text_encoder,
|
|
tokenizer=pipeline.tokenizer,
|
|
text_encoder_2=pipeline.text_encoder_2,
|
|
tokenizer_2=pipeline.tokenizer_2,
|
|
transformer=pipeline.transformer,
|
|
scheduler=pipeline.scheduler,
|
|
safety_checker=None,
|
|
feature_extractor=None,
|
|
controlnet=controlnets[0] if isinstance(controlnets, list) else controlnets, # can be a list
|
|
)
|
|
elif detect.is_zimage(pipeline) and len(controlnets) > 0:
|
|
from diffusers import ZImageControlNetPipeline
|
|
self.pipeline = ZImageControlNetPipeline(
|
|
vae=pipeline.vae,
|
|
text_encoder=pipeline.text_encoder,
|
|
tokenizer=pipeline.tokenizer,
|
|
transformer=pipeline.transformer,
|
|
scheduler=pipeline.scheduler,
|
|
controlnet=controlnets[0] if isinstance(controlnets, list) else controlnets, # can be a list
|
|
)
|
|
self.pipeline.task_args = { 'guidance_scale': 1 }
|
|
elif len(loras) > 0:
|
|
self.pipeline = pipeline
|
|
for lora in loras:
|
|
log.debug(f'Control {what} pipeline: lora="{lora}"')
|
|
lora = lora.replace('/huggingface.co/', '')
|
|
self.pipeline.load_lora_weights(lora)
|
|
"""
|
|
if p is not None:
|
|
p.prompt += f'<lora:{lora}:1.0>'
|
|
"""
|
|
else:
|
|
log.error(f'Control {what} pipeline: class={pipeline.__class__.__name__} unsupported model type')
|
|
return
|
|
|
|
if self.pipeline is None:
|
|
log.error(f'Control {what} pipeline: not initialized')
|
|
return
|
|
if dtype is not None:
|
|
self.pipeline = self.pipeline.to(dtype)
|
|
|
|
controlnet = None # free up memory
|
|
controlnets = None
|
|
sd_models.copy_diffuser_options(self.pipeline, pipeline)
|
|
if opts.diffusers_offload_mode == 'none':
|
|
sd_models.move_model(self.pipeline, devices.device)
|
|
sd_models.clear_caches()
|
|
sd_models.set_diffuser_offload(self.pipeline, 'model', force=True)
|
|
|
|
t1 = time.time()
|
|
debug_log(f'Control {what} pipeline: class={self.pipeline.__class__.__name__} time={t1-t0:.2f}')
|
|
|
|
def restore(self):
|
|
if self.pipeline is not None and hasattr(self.pipeline, 'unload_lora_weights'):
|
|
self.pipeline.unload_lora_weights()
|
|
self.pipeline = None
|
|
return self.orig_pipeline
|