diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e554bfd1..435df5506 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -123,6 +123,7 @@ For details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/master - installer dont cache installed state - fix Cosmos-Predict2 retrying TAESD download - better handle startup import errors + - fix traceback width preventing copy&paste - fix ansi controle output from scripts/extensions - fix diffusers models non-unique hash - fix loading of manually downloaded diffuser models diff --git a/installer.py b/installer.py index ae508d6d2..63232d35c 100644 --- a/installer.py +++ b/installer.py @@ -1,3 +1,4 @@ +from typing import List, Optional import os import sys import json @@ -104,6 +105,15 @@ def install_traceback(suppress: list = []): # setup console and file logging def setup_logging(): + from functools import partial, partialmethod + from logging.handlers import RotatingFileHandler + from rich.theme import Theme + from rich.logging import RichHandler + from rich.console import Console + from rich.padding import Padding + from rich.segment import Segment + from rich import print as rprint + from rich.pretty import install as pretty_install class RingBuffer(logging.StreamHandler): def __init__(self, capacity): @@ -128,6 +138,7 @@ def setup_logging(): def get(self): return self.buffer + class LogFilter(logging.Filter): def __init__(self): super().__init__() @@ -135,14 +146,35 @@ def setup_logging(): def filter(self, record): return len(record.getMessage()) > 2 + def override_padding(self, console, options): + style = console.get_style(self.style) + width = options.max_width + self.left = 0 + render_options = options.update_width(width - self.left - self.right) + if render_options.height is not None: + render_options = render_options.update_height(height=render_options.height - self.top - self.bottom) + lines = console.render_lines(self.renderable, render_options, style=style, pad=False) + _Segment = Segment + left = _Segment(" " * self.left, style) if self.left else None + right = [_Segment.line()] + blank_line: Optional[List[Segment]] = None + if self.top: + blank_line = [_Segment(f'{" " * width}\n', style)] + yield from blank_line * self.top + if left: + for line in lines: + yield left + yield from line + yield from right + else: + for line in lines: + yield from line + yield from right + if self.bottom: + blank_line = blank_line or [_Segment(f'{" " * width}\n', style)] + yield from blank_line * self.bottom + t_start = time.time() - from functools import partial, partialmethod - from logging.handlers import RotatingFileHandler - from rich.theme import Theme - from rich.logging import RichHandler - from rich.console import Console - from rich import print as rprint - from rich.pretty import install as pretty_install if args.log: global log_file # pylint: disable=global-statement @@ -159,12 +191,14 @@ def setup_logging(): global console # pylint: disable=global-statement theme = Theme({ "traceback.border": "black", - "traceback.border.syntax_error": "black", "inspect.value.border": "black", + "traceback.border.syntax_error": "dark_red", "logging.level.info": "blue_violet", "logging.level.debug": "purple4", "logging.level.trace": "dark_blue", }) + + Padding.__rich_console__ = override_padding console = Console( log_time=True, log_time_format='%H:%M:%S-%f', @@ -173,6 +207,7 @@ def setup_logging(): safe_box=True, theme=theme, ) + logging.basicConfig(level=logging.ERROR, format='%(asctime)s | %(name)s | %(levelname)s | %(module)s | %(message)s', handlers=[logging.NullHandler()]) # redirect default logger to null pretty_install(console=console) install_traceback() diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index 0457af7ea..cf3f5b5ee 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -130,7 +130,7 @@ def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str ) processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) if 'LLM' in shared.opts.cuda_compile: - model = compile_torch(model) + model = sd_models_compile.compile_torch(model) loaded = repo devices.torch_gc() sd_models.move_model(model, devices.device) @@ -180,7 +180,7 @@ def gemma(question: str, image: Image.Image, repo: str = None, system_prompt: st **quant_args, ) if 'LLM' in shared.opts.cuda_compile: - model = compile_torch(model) + model = sd_models_compile.compile_torch(model) processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) loaded = repo devices.torch_gc() @@ -307,7 +307,7 @@ def smol(question: str, image: Image.Image, repo: str = None, system_prompt: str ) processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) if 'LLM' in shared.opts.cuda_compile: - model = compile_torch(model) + model = sd_models_compile.compile_torch(model) loaded = repo devices.torch_gc() sd_models.move_model(model, devices.device) diff --git a/pipelines/model_chroma.py b/pipelines/model_chroma.py index a8c5bad42..f8750782b 100644 --- a/pipelines/model_chroma.py +++ b/pipelines/model_chroma.py @@ -118,20 +118,18 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unus } if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'): raise NotImplementedError('Nunchaku does not support Chroma Model yet. See https://github.com/mit-han-lab/nunchaku/issues/167') - elif 'transformer' not in kwargs and model_quant.check_quant('Model'): + if 'transformer' not in kwargs and model_quant.check_quant('Model'): load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True, modules_to_not_convert=["distilled_guidance_layer"]) - if quant_args: - kwargs['transformer'] = diffusers.ChromaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", **load_args, **quant_args) + kwargs['transformer'] = diffusers.ChromaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", **load_args, **quant_args) if 'text_encoder' not in kwargs and model_quant.check_nunchaku('TE'): import nunchaku nunchaku_precision = nunchaku.utils.get_precision() nunchaku_repo = 'mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors' shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}') kwargs['text_encoder'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype) - elif 'text_encoder' not in kwargs and model_quant.check_quant('TE'): + if 'text_encoder' not in kwargs and model_quant.check_quant('TE'): load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - if quant_args: - kwargs['text_encoder'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder", **load_args, **quant_args) + kwargs['text_encoder'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder", **load_args, **quant_args) except Exception as e: shared.log.error(f'Quantization: {e}') errors.display(e, 'Quantization:') diff --git a/pipelines/model_flux.py b/pipelines/model_flux.py index 77c0e768e..6ac58a707 100644 --- a/pipelines/model_flux.py +++ b/pipelines/model_flux.py @@ -132,10 +132,9 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unus kwargs['transformer'].quantization_method = 'SVDQuant' if shared.opts.nunchaku_attention: kwargs['transformer'].set_attention_impl("nunchaku-fp16") - elif 'transformer' not in kwargs and model_quant.check_quant('Model'): + if 'transformer' not in kwargs and model_quant.check_quant('Model'): load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - if quant_args: - kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", **load_args, **quant_args) + kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", **load_args, **quant_args) if 'text_encoder_2' not in kwargs and model_quant.check_nunchaku('TE'): import nunchaku nunchaku_precision = nunchaku.utils.get_precision() @@ -143,10 +142,9 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unus shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}') kwargs['text_encoder_2'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype) kwargs['text_encoder_2'].quantization_method = 'SVDQuant' - elif 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'): + if 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'): load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - if quant_args: - kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", **load_args, **quant_args) + kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", **load_args, **quant_args) except Exception as e: shared.log.error(f'Quantization: {e}') errors.display(e, 'Quantization:') diff --git a/pipelines/model_sana.py b/pipelines/model_sana.py index 7d0b65b62..c0fa486b7 100644 --- a/pipelines/model_sana.py +++ b/pipelines/model_sana.py @@ -15,12 +15,10 @@ def load_quants(kwargs, repo_id, cache_dir): kwargs['transformer'] = nunchaku.NunchakuSanaTransformer2DModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype) elif model_quant.check_quant('Model'): load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='Model') - if quant_args: - kwargs['transformer'] = diffusers.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args) + kwargs['transformer'] = diffusers.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args) if model_quant.check_quant('TE'): load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='TE') - if quant_args: - kwargs['text_encoder'] = transformers.AutoModelForCausalLM.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=cache_dir, **load_args, **quant_args) + kwargs['text_encoder'] = transformers.AutoModelForCausalLM.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=cache_dir, **load_args, **quant_args) return kwargs