mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
change onboarding and remove download default model
This commit is contained in:
Vladimir Mandic
@@ -0,0 +1,169 @@
|
||||
import os
|
||||
import time
|
||||
from typing import Union
|
||||
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionXLControlNetPipeline
|
||||
from modules.shared import log, opts
|
||||
from modules import errors
|
||||
|
||||
|
||||
what = 'ControlNet'
|
||||
debug = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
|
||||
debug('Trace: CONTROL')
|
||||
predefined_sd15 = {
|
||||
'OpenPose': "lllyasviel/control_v11p_sd15_openpose",
|
||||
'Canny': "lllyasviel/control_v11p_sd15_canny",
|
||||
'MLDS': "lllyasviel/control_v11p_sd15_mlsd",
|
||||
'Scribble': "lllyasviel/control_v11p_sd15_scribble",
|
||||
'SoftEdge': "lllyasviel/control_v11p_sd15_softedge",
|
||||
'Segment': "lllyasviel/control_v11p_sd15_seg",
|
||||
'Depth': "lllyasviel/control_v11f1p_sd15_depth",
|
||||
'NormalBae': "lllyasviel/control_v11p_sd15_normalbae",
|
||||
'LineArt': "lllyasviel/control_v11p_sd15_lineart",
|
||||
'LineArt Anime': "lllyasviel/control_v11p_sd15s2_lineart_anime",
|
||||
'Shuffle': "lllyasviel/control_v11e_sd15_shuffle",
|
||||
'IP2P': "lllyasviel/control_v11e_sd15_ip2p",
|
||||
'HED': "lllyasviel/sd-controlnet-hed",
|
||||
'Tile': "lllyasviel/control_v11f1e_sd15_tile",
|
||||
'TemporalNet': "CiaraRowles/TemporalNet",
|
||||
}
|
||||
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',
|
||||
}
|
||||
models = {}
|
||||
all_models = {}
|
||||
all_models.update(predefined_sd15)
|
||||
all_models.update(predefined_sdxl)
|
||||
cache_dir = 'models/control/controlnet'
|
||||
|
||||
|
||||
def find_models():
|
||||
path = os.path.join(opts.control_dir, 'controlnet')
|
||||
files = os.listdir(path)
|
||||
files = [f for f in files if f.endswith('.safetensors')]
|
||||
downloaded_models = {}
|
||||
for f in files:
|
||||
basename = os.path.splitext(f)[0]
|
||||
downloaded_models[basename] = os.path.join(path, f)
|
||||
all_models.update(downloaded_models)
|
||||
return downloaded_models
|
||||
|
||||
|
||||
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'] + sorted(predefined_sdxl) + sorted(find_models())
|
||||
elif modules.shared.sd_model_type == 'sd':
|
||||
models = ['None'] + sorted(predefined_sd15) + sorted(find_models())
|
||||
else:
|
||||
log.warning(f'Control {what} model list failed: unknown model type')
|
||||
models = ['None'] + sorted(predefined_sd15) + sorted(predefined_sdxl) + sorted(find_models())
|
||||
debug(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 model_id is not None:
|
||||
self.load()
|
||||
|
||||
def reset(self):
|
||||
if self.model is not None:
|
||||
log.debug(f'Control {what} model unloaded')
|
||||
self.model = None
|
||||
self.model_id = None
|
||||
|
||||
def load(self, model_id: str = None) -> str:
|
||||
try:
|
||||
t0 = time.time()
|
||||
model_id = model_id or self.model_id
|
||||
if model_id is None or model_id == 'None':
|
||||
self.reset()
|
||||
return
|
||||
model_path = all_models[model_id]
|
||||
if model_path == '':
|
||||
return
|
||||
if model_path is None:
|
||||
log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
|
||||
return
|
||||
log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}"')
|
||||
if model_path.endswith('.safetensors'):
|
||||
self.model = ControlNetModel.from_single_file(model_path, **self.load_config)
|
||||
else:
|
||||
self.model = ControlNetModel.from_pretrained(model_path, **self.load_config)
|
||||
if self.device is not None:
|
||||
self.model.to(self.device)
|
||||
if self.dtype is not None:
|
||||
self.model.to(self.dtype)
|
||||
t1 = time.time()
|
||||
self.model_id = model_id
|
||||
log.debug(f'Control {what} model loaded: id="{model_id}" path="{model_path}" time={t1-t0:.2f}')
|
||||
return f'{what} loaded model: {model_id}'
|
||||
except Exception as e:
|
||||
log.error(f'Control {what} model load failed: id="{model_id}" error={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], dtype = None):
|
||||
t0 = time.time()
|
||||
self.orig_pipeline = pipeline
|
||||
self.pipeline = None
|
||||
if pipeline is None:
|
||||
log.error('Control model pipeline: model not loaded')
|
||||
return
|
||||
elif isinstance(pipeline, StableDiffusionXLPipeline):
|
||||
self.pipeline = StableDiffusionXLControlNetPipeline(
|
||||
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),
|
||||
controlnet=controlnet, # can be a list
|
||||
).to(pipeline.device)
|
||||
elif isinstance(pipeline, StableDiffusionPipeline):
|
||||
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),
|
||||
requires_safety_checker=False,
|
||||
safety_checker=None,
|
||||
controlnet=controlnet, # can be a list
|
||||
).to(pipeline.device)
|
||||
else:
|
||||
log.error(f'Control {what} pipeline: class={pipeline.__class__.__name__} unsupported model type')
|
||||
return
|
||||
if dtype is not None and self.pipeline is not None:
|
||||
self.pipeline = self.pipeline.to(dtype)
|
||||
t1 = time.time()
|
||||
if self.pipeline is not None:
|
||||
log.debug(f'Control {what} pipeline: class={self.pipeline.__class__.__name__} time={t1-t0:.2f}')
|
||||
else:
|
||||
log.error(f'Control {what} pipeline: not initialized')
|
||||
|
||||
def restore(self):
|
||||
self.pipeline = None
|
||||
return self.orig_pipeline
|
||||
Reference in New Issue
Block a user