diff --git a/installer.py b/installer.py index ca6a690db..696e2474c 100644 --- a/installer.py +++ b/installer.py @@ -182,6 +182,16 @@ def package_version(package): return None +@lru_cache() +def package_spec(package): + spec = pkg_resources.working_set.by_key.get(package, None) # more reliable than importlib + if spec is None: + spec = pkg_resources.working_set.by_key.get(package.lower(), None) # check name variations + if spec is None: + spec = pkg_resources.working_set.by_key.get(package.replace('_', '-'), None) # check name variations + return spec + + # check if package is installed @lru_cache() def installed(package, friendly: str = None, reload = False, quiet = False): @@ -203,21 +213,17 @@ def installed(package, friendly: str = None, reload = False, quiet = False): p = pkg.split('>=') else: p = pkg.split('==') - spec = pkg_resources.working_set.by_key.get(p[0], None) # more reliable than importlib - if spec is None: - spec = pkg_resources.working_set.by_key.get(p[0].lower(), None) # check name variations - if spec is None: - spec = pkg_resources.working_set.by_key.get(p[0].replace('_', '-'), None) # check name variations + spec = package_spec(p[0]) ok = ok and spec is not None if ok: - package_version = pkg_resources.get_distribution(p[0]).version + pkg_version = package_version(p[0]) if len(p) > 1: - exact = package_version == p[1] + exact = pkg_version == p[1] if not exact and not quiet: if args.experimental: - log.warning(f"Package: {p[0]} {package_version} required {p[1]} allowing experimental") + log.warning(f"Package: {p[0]} {pkg_version} required {p[1]} allowing experimental") else: - log.warning(f"Package: {p[0]} {package_version} required {p[1]} version mismatch") + log.warning(f"Package: {p[0]} {pkg_version} required {p[1]} version mismatch") ok = ok and (exact or args.experimental) else: if not quiet: diff --git a/modules/ctrlx/__init__.py b/modules/ctrlx/__init__.py index c5afc2926..caa0e44a1 100644 --- a/modules/ctrlx/__init__.py +++ b/modules/ctrlx/__init__.py @@ -56,7 +56,6 @@ class CtrlXStableDiffusionXLPipeline(StableDiffusionXLPipeline): # diffusers==0 dtype, device, generator=None, noise=None, ): batch_size = batch_size * num_images_per_prompt - if noise is None: shape = ( batch_size, diff --git a/modules/processing_args.py b/modules/processing_args.py index 1aebca869..2d6fa58d3 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -99,7 +99,7 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 steps = kwargs.get("num_inference_steps", None) or len(getattr(p, 'timesteps', ['1'])) clip_skip = kwargs.pop("clip_skip", 1) - prompt_parser_diffusers.fix_position_ids(model) + # prompt_parser_diffusers.fix_position_ids(model) if shared.opts.prompt_attention != 'Fixed attention' and 'Onnx' not in model.__class__.__name__ and ( 'StableDiffusion' in model.__class__.__name__ or 'StableCascade' in model.__class__.__name__ or diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 1cb0c936b..3a8f5d2c8 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -13,6 +13,7 @@ from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, c debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None debug('Trace: DIFFUSERS') last_p = None +orig_pipeline = shared.sd_model def restore_state(p: processing.StableDiffusionProcessing): @@ -334,7 +335,6 @@ def process_decode(p: processing.StableDiffusionProcessing, output): return results -orig_pipeline = shared.sd_model def update_pipeline(sd_model, p: processing.StableDiffusionProcessing): if sd_models.get_diffusers_task(sd_model) == sd_models.DiffusersTaskType.INPAINTING and getattr(p, 'image_mask', None) is None and p.task_args.get('image_mask', None) is None and getattr(p, 'mask', None) is None: shared.log.warning('Processing: mode=inpaint mask=None') @@ -355,6 +355,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing): debug(f'Process diffusers args: {vars(p)}') results = [] p = restore_state(p) + global orig_pipeline # pylint: disable=global-statement + orig_pipeline = shared.sd_model if shared.state.interrupted or shared.state.skipped: shared.sd_model = orig_pipeline diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 75e9f5b32..c063833a2 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -5,7 +5,7 @@ import typing import torch from compel.embeddings_provider import BaseTextualInversionManager, EmbeddingsProvider from transformers import PreTrainedTokenizer -from modules import shared, prompt_parser, devices, sd_models, errors +from modules import shared, prompt_parser, devices, sd_models from modules.prompt_parser_xhinker import get_weighted_text_embeddings_sd15, get_weighted_text_embeddings_sdxl_2p, get_weighted_text_embeddings_sd3, get_weighted_text_embeddings_flux1 debug_enabled = os.environ.get('SD_PROMPT_DEBUG', None) @@ -17,29 +17,9 @@ token_type = None # used by helper get_tokens cache = {} -def fix_position_ids(pipe): - # position_ids are created on te creation and are simple index cache - # but somehow can be corrupt in CLIPTextEmbeddings forward call - # see transformers/models/clip/modeling_clip.py:CLIPTextEmbeddings - # reproduction: load sdxl model -> generate -> generate -> load sdxl model -> generate -> generate -> load sdxl model -> generate -> generate - if hasattr(pipe, 'text_encoder') and pipe.text_encoder.text_model.embeddings.position_ids[0][0] > 0: - debug(f'TE1 fix: ids={pipe.text_encoder.text_model.embeddings.position_ids}') - pipe.text_encoder.text_model.embeddings.position_ids = torch.arange(pipe.text_encoder.config.max_position_embeddings).expand((1, -1)).to(pipe.text_encoder.device) - if hasattr(pipe, 'text_encoder_2') and pipe.text_encoder_2.text_model.embeddings.position_ids[0][0] > 0: - debug(f'TE2 fix: ids={pipe.text_encoder_2.text_model.embeddings.position_ids}') - pipe.text_encoder_2.text_model.embeddings.position_ids = torch.arange(pipe.text_encoder_2.config.max_position_embeddings).expand((1, -1)).to(pipe.text_encoder_2.device) - - -def compel_hijack(self, token_ids: torch.Tensor, - attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor: +def compel_hijack(self, token_ids: torch.Tensor, attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor: needs_hidden_states = self.returned_embeddings_type != 1 - try: # can crash in ATen/native/cuda/Indexing since position_ids are corrupt so index lookup fails, but its not compel specific, happens with fixed attention as well - sd_models.move_model(self.text_encoder, devices.device) - text_encoder_output = self.text_encoder(token_ids, attention_mask, output_hidden_states=needs_hidden_states, return_dict=True) - except Exception as e: # its a non-recoverable error as cuda state is corrupt - shared.log.error(f'TE: class={self.text_encoder.__class__} device={self.text_encoder.device} dtype={self.text_encoder.dtype} {e}') - errors.display(e, 'TE:') - return None + text_encoder_output = self.text_encoder(token_ids, attention_mask, output_hidden_states=needs_hidden_states, return_dict=True) if not needs_hidden_states: return text_encoder_output.last_hidden_state @@ -60,15 +40,14 @@ def compel_hijack(self, token_ids: torch.Tensor, return hidden_state -def sd3_compel_hijack(self, token_ids: torch.Tensor, - attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor: +def sd3_compel_hijack(self, token_ids: torch.Tensor, attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor: needs_hidden_states = True text_encoder_output = self.text_encoder(token_ids, attention_mask, output_hidden_states=needs_hidden_states, return_dict=True) clip_skip = int(self.returned_embeddings_type) hidden_state = text_encoder_output.hidden_states[-(clip_skip+1)] - return hidden_state + def insert_parser_highjack(pipename): if "StableDiffusion3" in pipename: EmbeddingsProvider._encode_token_ids_to_embeddings = sd3_compel_hijack # pylint: disable=protected-access @@ -192,14 +171,9 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c pipe.maybe_free_model_hooks() devices.torch_gc() - p.prompt_embeds = [] - p.positive_pooleds = [] - p.negative_embeds = [] - p.negative_pooleds = [] - p.scheduled_prompt = False + prompt_embeds, positive_pooleds, negative_embeds, negative_pooleds = [], [], [], [] last_prompt, last_negative = None, None for prompt, negative in zip(prompts, negative_prompts): - prompt_embeds, positive_pooleds, negative_embeds, negative_pooleds = [], [], [], [] prompt_embed, positive_pooled, negative_embed, negative_pooled = None, None, None, None if last_prompt == prompt and last_negative == negative: prompt_embeds.append(prompt_embeds[-1]) @@ -211,7 +185,11 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c continue positive_schedule, scheduled = get_prompt_schedule(prompt, steps) negative_schedule, neg_scheduled = get_prompt_schedule(negative, steps) - p.scheduled_prompt = p.scheduled_prompt or scheduled or neg_scheduled + p.scheduled_prompt = scheduled or neg_scheduled + p.prompt_embeds = [] + p.positive_pooleds = [] + p.negative_embeds = [] + p.negative_pooleds = [] for i in range(max(len(positive_schedule), len(negative_schedule))): positive_prompt = positive_schedule[i % len(positive_schedule)] @@ -230,25 +208,25 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c negative_pooleds.append(negative_pooled) last_prompt, last_negative = prompt, negative - def fix_length(embeds): - max_len = max([e.shape[1] for e in embeds if e is not None]) - for i, e in enumerate(embeds): - if e is not None and e.shape[1] < max_len: - expanded = torch.zeros((e.shape[0], max_len, e.shape[2]), device=e.device, dtype=e.dtype) - expanded[:, :e.shape[1], :] = e - embeds[i] = expanded - return torch.cat(embeds, dim=0).to(devices.device, dtype=devices.dtype) + def fix_length(embeds): + max_len = max([e.shape[1] for e in embeds if e is not None]) + for i, e in enumerate(embeds): + if e is not None and e.shape[1] < max_len: + expanded = torch.zeros((e.shape[0], max_len, e.shape[2]), device=e.device, dtype=e.dtype) + expanded[:, :e.shape[1], :] = e + embeds[i] = expanded + return torch.cat(embeds, dim=0).to(devices.device, dtype=devices.dtype) - if len(prompt_embeds) > 0: - p.prompt_embeds.append(fix_length(prompt_embeds)) - if len(negative_embeds) > 0: - p.negative_embeds.append(fix_length(negative_embeds)) - if len(positive_pooleds) > 0: - p.positive_pooleds.append(fix_length(positive_pooleds)) - if len(negative_pooleds) > 0: - p.negative_pooleds.append(fix_length(negative_pooleds)) + if len(prompt_embeds) > 0: + p.prompt_embeds.append(fix_length(prompt_embeds)) + if len(negative_embeds) > 0: + p.negative_embeds.append(fix_length(negative_embeds)) + if len(positive_pooleds) > 0: + p.positive_pooleds.append(fix_length(positive_pooleds)) + if len(negative_pooleds) > 0: + p.negative_pooleds.append(fix_length(negative_pooleds)) - if p.batch_size == 1: + if shared.opts.sd_textencoder_cache and p.batch_size == 1: cache.update({ 'prompt_embeds': p.prompt_embeds, 'negative_embeds': p.negative_embeds, diff --git a/modules/sd_models.py b/modules/sd_models.py index 577a74ff5..9e1d3c270 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -737,6 +737,7 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True): if offload: set_diffuser_offload(sd_model, op) + def set_diffuser_offload(sd_model, op: str = 'model'): if not shared.native: shared.log.warning('Attempting to use offload with backend=original') @@ -746,43 +747,45 @@ def set_diffuser_offload(sd_model, op: str = 'model'): return if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False - if hasattr(sd_model, "enable_model_cpu_offload"): - if shared.opts.diffusers_offload_mode == "model": - try: - shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') - if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: - shared.opts.diffusers_move_base = False - shared.opts.diffusers_move_unet = False - shared.opts.diffusers_move_refiner = False - shared.log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled') - if not hasattr(sd_model, "_all_hooks") or len(sd_model._all_hooks) == 0: # pylint: disable=protected-access - sd_model.enable_model_cpu_offload(device=devices.device) + if hasattr(sd_model, 'maybe_free_model_hooks') and shared.opts.diffusers_offload_mode == "none": + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') + sd_model.maybe_free_model_hooks() + sd_model.has_accelerate = False + if hasattr(sd_model, "enable_model_cpu_offload") and shared.opts.diffusers_offload_mode == "model": + try: + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') + if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: + shared.opts.diffusers_move_base = False + shared.opts.diffusers_move_unet = False + shared.opts.diffusers_move_refiner = False + shared.log.warning(f'Disabling {op} "Move model to CPU" since "Model CPU offload" is enabled') + if not hasattr(sd_model, "_all_hooks") or len(sd_model._all_hooks) == 0: # pylint: disable=protected-access + sd_model.enable_model_cpu_offload(device=devices.device) + else: + sd_model.maybe_free_model_hooks() + sd_model.has_accelerate = True + except Exception as e: + shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') + if hasattr(sd_model, "enable_sequential_cpu_offload") and shared.opts.diffusers_offload_mode == "sequential": + try: + shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') + if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: + shared.opts.diffusers_move_base = False + shared.opts.diffusers_move_unet = False + shared.opts.diffusers_move_refiner = False + shared.log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled') + if sd_model.has_accelerate: + if op == "vae": # reapply sequential offload to vae + from accelerate import cpu_offload + sd_model.vae.to("cpu") + cpu_offload(sd_model.vae, devices.device, offload_buffers=len(sd_model.vae._parameters) > 0) # pylint: disable=protected-access else: - sd_model.maybe_free_model_hooks() - sd_model.has_accelerate = True - except Exception as e: - shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') - if hasattr(sd_model, "enable_sequential_cpu_offload"): - if shared.opts.diffusers_offload_mode == "sequential": - try: - shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') - if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: - shared.opts.diffusers_move_base = False - shared.opts.diffusers_move_unet = False - shared.opts.diffusers_move_refiner = False - shared.log.warning(f'Disabling {op} "Move model to CPU" since "Sequential CPU offload" is enabled') - if sd_model.has_accelerate: - if op == "vae": # reapply sequential offload to vae - from accelerate import cpu_offload - sd_model.vae.to("cpu") - cpu_offload(sd_model.vae, devices.device, offload_buffers=len(sd_model.vae._parameters) > 0) # pylint: disable=protected-access - else: - pass # do nothing if offload is already applied - else: - sd_model.enable_sequential_cpu_offload(device=devices.device) - sd_model.has_accelerate = True - except Exception as e: - shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') + pass # do nothing if offload is already applied + else: + sd_model.enable_sequential_cpu_offload(device=devices.device) + sd_model.has_accelerate = True + except Exception as e: + shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}') if shared.opts.diffusers_offload_mode == "balanced": try: shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode}') diff --git a/requirements.txt b/requirements.txt index 5433f4061..5a546acd7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -40,7 +40,7 @@ clip-interrogator==0.6.0 antlr4-python3-runtime==4.9.3 requests==2.32.3 tqdm==4.66.5 -accelerate==0.34.2 +accelerate==0.33.0 opencv-contrib-python-headless==4.9.0.80 einops==0.4.1 gradio==3.43.2