diff --git a/modules/hashes.py b/modules/hashes.py index 84071bfb5..288a175d5 100644 --- a/modules/hashes.py +++ b/modules/hashes.py @@ -25,7 +25,7 @@ def calculate_sha256(filename, quiet=False): hash_sha256 = hashlib.sha256() blksize = 1024 * 1024 if not quiet: - with progress.open(filename, 'rb', description=f'Calculating model hash: [cyan]{filename}', auto_refresh=True) as f: + with progress.open(filename, 'rb', description=f'Calculating model hash: [cyan]{filename}', auto_refresh=True, console=shared.console) as f: for chunk in iter(lambda: f.read(blksize), b""): hash_sha256.update(chunk) else: @@ -57,7 +57,7 @@ def sha256(filename, title, use_addnet_hash=False): if not os.path.isfile(filename): return None if use_addnet_hash: - with progress.open(filename, 'rb', description=f'Calculating model hash: [cyan]{filename}', auto_refresh=True) as f: + with progress.open(filename, 'rb', description=f'Calculating model hash: [cyan]{filename}', auto_refresh=True, console=shared.console) as f: sha256_value = addnet_hash_safetensors(f) else: sha256_value = calculate_sha256(filename) diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py index 351e261ea..d9cc95e86 100644 --- a/modules/hypernetworks/hypernetwork.py +++ b/modules/hypernetworks/hypernetwork.py @@ -224,7 +224,7 @@ class Hypernetwork: self.filename = filename if self.name is None: self.name = os.path.splitext(os.path.basename(filename))[0] - with progress.open(filename, 'rb', description=f'Loading hypernetwork: [cyan]{filename}', auto_refresh=True) as f: + with progress.open(filename, 'rb', description=f'Loading hypernetwork: [cyan]{filename}', auto_refresh=True, console=shared.console) as f: state_dict = torch.load(f, map_location='cpu') self.layer_structure = state_dict.get('layer_structure', [1, 2, 1]) self.optional_info = state_dict.get('optional_info', None) diff --git a/modules/modelloader.py b/modules/modelloader.py index 40136479d..f9b6048a0 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -70,7 +70,7 @@ def download_civit_preview(model_path: str, preview_url: str): shared.state.begin('civitai-download-preview') try: with open(preview_file, 'wb') as f: - with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn()) as progress: + with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), console=shared.console) as progress: task = progress.add_task(description="Download starting", total=total_size) for data in req.iter_content(block_size): written = written + len(data) @@ -110,7 +110,7 @@ def download_civit_model(model_url: str, model_name: str, model_path: str, model shared.state.begin('civitai-download-model') try: with open(model_file, 'wb') as f: - with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn()) as progress: + with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), console=shared.console) as progress: task = progress.add_task(description="Download starting", total=total_size) # for data in tqdm(req.iter_content(block_size), total=total_size//1024, unit='KB', unit_scale=False): for data in req.iter_content(block_size): diff --git a/modules/models/diffusion/uni_pc/uni_pc.py b/modules/models/diffusion/uni_pc/uni_pc.py index 56604a7b6..fa16a0b48 100644 --- a/modules/models/diffusion/uni_pc/uni_pc.py +++ b/modules/models/diffusion/uni_pc/uni_pc.py @@ -757,7 +757,7 @@ class UniPC: #print(f"Running UniPC Sampling with {timesteps.shape[0]} timesteps, order {order}") assert steps >= order, "UniPC order must be < sampling steps" assert timesteps.shape[0] - 1 == steps - with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn()) as progress: + with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), console=shared.console) as progress: task = progress.add_task(description="Initializing", total=steps) t = time.time() with devices.inference_context(): diff --git a/modules/paths.py b/modules/paths.py index c0d7cbd11..9530ff5ba 100644 --- a/modules/paths.py +++ b/modules/paths.py @@ -62,7 +62,7 @@ def create_paths(opts, log=None): def fix_path(folder): tgt = opts.data.get(folder, None) or opts.data_labels[folder].default if tgt is None or tgt == '': - return + return tgt if len(data_path) > 0 and tgt.startswith(data_path): # path is already relative to data_path return tgt fullpath = os.path.join(data_path, tgt) diff --git a/modules/script_loading.py b/modules/script_loading.py index 64f16e681..61f08527a 100644 --- a/modules/script_loading.py +++ b/modules/script_loading.py @@ -13,8 +13,8 @@ def load_module(path): module_spec = importlib.util.spec_from_file_location(os.path.basename(path), path) module = importlib.util.module_from_spec(module_spec) try: - stdout = io.StringIO() - with contextlib.redirect_stdout(stdout): + # stdout = io.StringIO() + with contextlib.redirect_stdout(io.StringIO()) as stdout: module_spec.loader.exec_module(module) setup_logging() # reset since scripts can hijaack logging for line in stdout.getvalue().splitlines(): diff --git a/modules/sd_models.py b/modules/sd_models.py index 2147e31db..8fba59aa1 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -387,7 +387,7 @@ def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unuse return None try: pl_sd = None - with progress.open(checkpoint_file, 'rb', description=f'[cyan]Loading weights: [yellow]{checkpoint_file}', auto_refresh=True) as f: + with progress.open(checkpoint_file, 'rb', description=f'[cyan]Loading weights: [yellow]{checkpoint_file}', auto_refresh=True, console=shared.console) as f: _, extension = os.path.splitext(checkpoint_file) if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt: shared.log.warning(f"Checkpoint loading disabled: {checkpoint_file}") diff --git a/modules/shared.py b/modules/shared.py index d98bbf448..b5fbbe12a 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -11,6 +11,7 @@ from enum import Enum import gradio as gr import tqdm import fasteners +from rich.console import Console from modules import errors, ui_components, shared_items, cmd_args from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack @@ -69,7 +70,7 @@ restricted_opts = { "outdir_init_images" } compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order'] - +console = Console(log_time=True, log_time_format='%H:%M:%S-%f') def is_url(string): parsed_url = urlparse(string) diff --git a/modules/styles.py b/modules/styles.py index 94052d548..abeeb5e08 100644 --- a/modules/styles.py +++ b/modules/styles.py @@ -41,6 +41,7 @@ class StyleDatabase: def __init__(self, opts): self.no_style = Style("None") self.styles = {} + self.path = opts.styles_dir if os.path.isfile(opts.styles_dir) or opts.styles_dir.endswith(".csv"): legacy_file = opts.styles_dir self.load_csv(legacy_file)