From 899c896f23871bc8ac2295e4f6614aa8f7487374 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 17 Oct 2024 00:25:03 +0300 Subject: [PATCH] cleanup --- modules/model_stablecascade.py | 9 +++------ modules/sd_models.py | 12 +++++------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/modules/model_stablecascade.py b/modules/model_stablecascade.py index 74158b8ca..44cea1818 100644 --- a/modules/model_stablecascade.py +++ b/modules/model_stablecascade.py @@ -132,11 +132,6 @@ def load_cascade_combined(checkpoint_info, diffusers_load_config): else: sd_model = StableCascadeCombinedPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) - shared.log.debug(f'StableCascade combined: {sd_model.__class__.__name__}') - return sd_model - - -def cascade_post_load(sd_model): sd_model.prior_pipe.scheduler.config.clip_sample = False sd_model.default_scheduler = copy.deepcopy(sd_model.prior_pipe.scheduler) sd_model.prior_pipe.get_timestep_ratio_conditioning = get_timestep_ratio_conditioning @@ -162,10 +157,12 @@ def cascade_post_load(sd_model): text_encoder=None, latent_dim_scale=sd_model.decoder_pipe.config.latent_dim_scale, ) + + shared.log.debug(f'StableCascade combined: {sd_model.__class__.__name__}') return sd_model -# Custom sampler support. Remove after the changes gets upstreamed: https://github.com/huggingface/diffusers/pull/9132 +# Balanced offload hooks: class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline): def guidance_scale(self): return self._guidance_scale diff --git a/modules/sd_models.py b/modules/sd_models.py index a4828189b..abf65e362 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -611,7 +611,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False): guess = 'PixArt-Alpha' if 'stable-diffusion-3' in f.lower(): guess = 'Stable Diffusion 3' - if 'stable-cascade' in f.lower() or 'stablecascade' in f.lower() or 'wuerstchen3' in f.lower() or 'sotediffusion' in f.lower(): + if 'stable-cascade' in f.lower() or 'stablecascade' in f.lower() or 'wuerstchen3' in f.lower() or ('sotediffusion' in f.lower() and "v2" in f.lower()): if devices.dtype == torch.float16: warn('Stable Cascade does not support Float16') guess = 'Stable Cascade' @@ -853,7 +853,7 @@ def apply_balanced_offload(sd_model): def apply_balanced_offload_to_module(pipe): for module_name in pipe._internal_dict.keys(): # pylint: disable=protected-access - module = getattr(pipe, module_name) + module = getattr(pipe, module_name, None) if isinstance(module, torch.nn.Module): checkpoint_name = pipe.sd_checkpoint_info.name if getattr(pipe, "sd_checkpoint_info", None) is not None else None if checkpoint_name is None: @@ -1058,9 +1058,8 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op=' sd_model = None try: if model_type in ['Stable Cascade']: # forced pipeline - from modules.model_stablecascade import load_cascade_combined, cascade_post_load + from modules.model_stablecascade import load_cascade_combined sd_model = load_cascade_combined(checkpoint_info, diffusers_load_config) - cascade_post_load(sd_model) elif model_type in ['InstaFlow']: # forced pipeline pipeline = diffusers.utils.get_class_from_dynamic_module('instaflow_one_step', module_file='pipeline.py') sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) @@ -1856,9 +1855,8 @@ def disable_offload(sd_model): if not getattr(sd_model, 'has_accelerate', False): return for _name, model in sd_model.components.items(): - if not isinstance(model, torch.nn.Module): - continue - remove_hook_from_module(model, recurse=True) + if isinstance(model, torch.nn.Module): + remove_hook_from_module(model, recurse=True) sd_model.has_accelerate = False