mirror of https://github.com/vladmandic/automatic
247 lines
8.9 KiB
Python
247 lines
8.9 KiB
Python
from __future__ import annotations
|
|
import io
|
|
import os
|
|
import sys
|
|
import time
|
|
import contextlib
|
|
|
|
from enum import Enum
|
|
from typing import TYPE_CHECKING
|
|
import gradio as gr
|
|
from installer import print_dict # pylint: disable=unused-import
|
|
from modules.logger import log
|
|
log.debug('Initializing: shared module')
|
|
|
|
import modules.memmon
|
|
import modules.paths as paths
|
|
from modules.json_helpers import readfile # pylint: disable=W0611
|
|
from modules.shared_helpers import listdir, req # pylint: disable=W0611
|
|
from modules import errors, devices, shared_state, cmd_args, theme, history, files_cache # pylint: disable=unused-import
|
|
from modules.shared_defaults import get_default_modes
|
|
from modules.memstats import memory_stats # pylint: disable=unused-import
|
|
|
|
log.debug('Initializing: pipelines')
|
|
from modules import shared_items # pylint: disable=unused-import
|
|
from modules.caption.openclip import get_clip_models, refresh_clip_models # pylint: disable=unused-import
|
|
from modules.caption.vqa import vlm_models, vlm_prompts, vlm_system, vlm_default # pylint: disable=unused-import
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
# Behavior modified by __future__.annotations
|
|
from diffusers import DiffusionPipeline
|
|
from modules.shared_legacy import LegacyOption
|
|
from modules.ui_extra_networks import ExtraNetworksPage
|
|
|
|
|
|
class Backend(Enum):
|
|
ORIGINAL = 1
|
|
DIFFUSERS = 2
|
|
|
|
|
|
errors.install([gr])
|
|
demo: gr.Blocks | None = None
|
|
api = None
|
|
url = 'https://github.com/vladmandic/sdnext'
|
|
cmd_opts = cmd_args.parse_args()
|
|
parser = cmd_args.parser
|
|
hide_dirs = {"visible": not cmd_opts.hide_ui_dir_config}
|
|
listfiles = listdir
|
|
xformers_available = False
|
|
compiled_model_state = None
|
|
sd_upscalers = []
|
|
detailers = []
|
|
yolo = None
|
|
tab_names = []
|
|
extra_networks: list[ExtraNetworksPage] = []
|
|
hypernetworks = {}
|
|
settings_components = {}
|
|
restricted_opts = {
|
|
"samples_filename_pattern",
|
|
"directories_filename_pattern",
|
|
"outdir_samples",
|
|
"outdir_txt2img_samples",
|
|
"outdir_img2img_samples",
|
|
"outdir_extras_samples",
|
|
"outdir_control_samples",
|
|
"outdir_grids",
|
|
"outdir_txt2img_grids",
|
|
"outdir_save",
|
|
"outdir_init_images"
|
|
}
|
|
resize_modes = ["None", "Fixed", "Crop", "Fill", "Outpaint", "Context aware"]
|
|
max_workers = 12
|
|
default_hfcache_dir = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(paths.models_path, 'huggingface')
|
|
state = shared_state.State()
|
|
models_path = paths.models_path
|
|
script_path = paths.script_path
|
|
data_path = paths.data_path
|
|
|
|
|
|
# early select backend
|
|
backend = Backend.DIFFUSERS
|
|
if cmd_opts.use_openvino: # override for openvino
|
|
os.environ.setdefault('PYTORCH_TRACING_MODE', 'TORCHFX')
|
|
from modules.intel.openvino import get_device_list as get_openvino_device_list # pylint: disable=ungrouped-imports,unused-import
|
|
elif cmd_opts.use_ipex or devices.has_xpu():
|
|
from modules.intel.ipex import ipex_init
|
|
ok, e = ipex_init()
|
|
if not ok:
|
|
log.error(f'IPEX initialization failed: {e}')
|
|
if os.environ.get('SD_DEVICE_DEBUG', None) is not None:
|
|
errors.display(e, 'IPEX')
|
|
elif cmd_opts.use_directml:
|
|
from modules.dml import directml_init
|
|
ok, e = directml_init()
|
|
if not ok:
|
|
log.error(f'DirectML initialization failed: {e}')
|
|
if os.environ.get('SD_DEVICE_DEBUG', None) is not None:
|
|
errors.display(e, 'DirectML')
|
|
elif cmd_opts.use_rocm or devices.has_rocm():
|
|
from modules.rocm import rocm_init
|
|
ok, e = rocm_init()
|
|
if not ok:
|
|
log.error(f'ROCm initialization failed: {e}')
|
|
if os.environ.get('SD_DEVICE_DEBUG', None) is not None:
|
|
errors.display(e, 'ROCm')
|
|
devices.backend = devices.get_backend(cmd_opts)
|
|
devices.device = devices.get_optimal_device()
|
|
mem_stat = memory_stats()
|
|
cpu_memory = round(mem_stat['ram']['total'] if "ram" in mem_stat else 0)
|
|
gpu_memory = round(mem_stat['gpu']['total'] if "gpu" in mem_stat else 0)
|
|
if gpu_memory == 0:
|
|
gpu_memory = cpu_memory
|
|
native = backend == Backend.DIFFUSERS
|
|
if not files_cache.do_cache_folders:
|
|
log.warning('File cache disabled: ')
|
|
|
|
|
|
def list_checkpoint_titles():
|
|
import modules.sd_models # pylint: disable=W0621
|
|
return modules.sd_models.checkpoint_titles()
|
|
|
|
|
|
list_checkpoint_tiles = list_checkpoint_titles # alias for legacy typo
|
|
default_sd_model_file = paths.default_sd_model_file
|
|
default_checkpoint = list_checkpoint_titles()[0] if len(list_checkpoint_titles()) > 0 else "model.safetensors"
|
|
|
|
|
|
def is_url(string):
|
|
from urllib.parse import urlparse
|
|
parsed_url = urlparse(string)
|
|
return all([parsed_url.scheme, parsed_url.netloc])
|
|
|
|
|
|
def refresh_checkpoints():
|
|
import modules.sd_models # pylint: disable=W0621
|
|
return modules.sd_models.list_models()
|
|
|
|
|
|
def refresh_vaes():
|
|
import modules.sd_vae # pylint: disable=W0621
|
|
modules.sd_vae.refresh_vae_list()
|
|
|
|
|
|
def refresh_upscalers():
|
|
import modules.modelloader # pylint: disable=W0621
|
|
modules.modelloader.load_upscalers()
|
|
|
|
|
|
def list_samplers():
|
|
import modules.sd_samplers # pylint: disable=W0621
|
|
modules.sd_samplers.set_samplers()
|
|
return modules.sd_samplers.all_samplers
|
|
|
|
log.debug('Initializing: default modes')
|
|
startup_offload_mode, startup_offload_min_gpu, startup_offload_max_gpu, startup_cross_attention, startup_sdp_options, startup_sdp_choices, startup_sdp_override_options, startup_sdp_override_choices, startup_offload_always, startup_offload_never = get_default_modes(cmd_opts=cmd_opts, mem_stat=mem_stat)
|
|
|
|
log.debug('Initializing: settings')
|
|
from modules import ui_definitions
|
|
from modules.ui_definitions import OptionInfo, options_section # pylint: disable=unused-import
|
|
options_templates = ui_definitions.create_settings(cmd_opts)
|
|
from modules.shared_legacy import get_legacy_options
|
|
options_templates.update(get_legacy_options())
|
|
from modules.options_handler import Options
|
|
config_filename = cmd_opts.config
|
|
secrets_filename = cmd_opts.secrets
|
|
opts = Options(options_templates, restricted=restricted_opts, filename=config_filename, secrets=secrets_filename)
|
|
cmd_opts = cmd_args.settings_args(opts, cmd_opts)
|
|
cmd_opts = cmd_args.override_args(opts, cmd_opts)
|
|
if cmd_opts.locale is not None:
|
|
opts.data['ui_locale'] = cmd_opts.locale
|
|
|
|
log.debug('Initializing: backend')
|
|
from modules.dml import directml_do_hijack
|
|
if cmd_opts.use_xformers:
|
|
opts.data['cross_attention_optimization'] = 'xFormers'
|
|
opts.data['uni_pc_lower_order_final'] = opts.schedulers_use_loworder # compatibility
|
|
opts.data['uni_pc_order'] = max(2, opts.schedulers_solver_order) # compatibility
|
|
log.info(f'Engine: backend={backend} compute={devices.backend} device={devices.get_optimal_device_name()} attention="{opts.cross_attention_optimization}" mode={devices.inference_context.__name__}')
|
|
|
|
profiler = None
|
|
import modules.styles
|
|
prompt_styles = modules.styles.StyleDatabase(opts)
|
|
reference_models = readfile(os.path.join('data', 'reference.json'), as_type="dict") if opts.extra_network_reference_enable else {}
|
|
cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or (cmd_opts.server_name or False)) and not cmd_opts.insecure
|
|
|
|
log.debug('Initializing: devices')
|
|
devices.args = cmd_opts
|
|
devices.opts = opts
|
|
devices.onnx = [opts.onnx_execution_provider]
|
|
devices.set_cuda_params()
|
|
if opts.onnx_cpu_fallback and 'CPUExecutionProvider' not in devices.onnx:
|
|
devices.onnx.append('CPUExecutionProvider')
|
|
device = devices.device
|
|
parallel_processing_allowed = not cmd_opts.lowvram
|
|
mem_mon = modules.memmon.MemUsageMonitor("MemMon", devices.device)
|
|
history = history.History()
|
|
if devices.backend == "directml":
|
|
directml_do_hijack()
|
|
log.debug('Quantization: registered=SDNQ')
|
|
|
|
try:
|
|
log.info(f'Device: {print_dict(devices.get_gpu_info())}')
|
|
except Exception as ex:
|
|
log.error(f'Device: {ex}')
|
|
|
|
|
|
def restart_server(restart=True):
|
|
if demo is None:
|
|
return
|
|
log.critical('Server shutdown requested')
|
|
try:
|
|
sys.tracebacklimit = 0
|
|
stdout = io.StringIO()
|
|
stderr = io.StringIO()
|
|
with contextlib.redirect_stdout(stdout), contextlib.redirect_stdout(stderr):
|
|
demo.server.wants_restart = restart
|
|
demo.server.should_exit = True
|
|
demo.server.force_exit = True
|
|
demo.close(verbose=False)
|
|
demo.server.close()
|
|
demo.fns = []
|
|
time.sleep(1)
|
|
sys.tracebacklimit = 100
|
|
# os._exit(0)
|
|
except (Exception, BaseException) as err:
|
|
log.error(f'Server shutdown error: {err}')
|
|
if restart:
|
|
log.info('Server will restart')
|
|
|
|
|
|
def restore_defaults(restart=True):
|
|
if os.path.exists(cmd_opts.config):
|
|
log.info('Restoring server defaults')
|
|
os.remove(cmd_opts.config)
|
|
restart_server(restart)
|
|
|
|
|
|
# startup def of shared.sd_model before its redefined in modeldata
|
|
sd_model: DiffusionPipeline | None = None # dummy and overwritten by class
|
|
sd_refiner: DiffusionPipeline | None = None # dummy and overwritten by class
|
|
sd_model_type: str = '' # dummy and overwritten by class
|
|
sd_refiner_type: str = '' # dummy and overwritten by class
|
|
sd_loaded: bool = False # dummy and overwritten by class
|
|
|
|
from modules.modeldata import Shared # pylint: disable=ungrouped-imports
|
|
sys.modules[__name__].__class__ = Shared
|