cleanup

Signed-off-by: vladmandic <mandic00@live.com>

modules/sd_checkpoint.py

@@ -6,7 +6,7 @@ import time
 import json
 import collections
 from PIL import Image
-from modules import shared, paths, modelloader, hashes, sd_hijack_accelerate
+from modules import shared, paths, modelloader, hashes
 from modules.logger import log
 from modules.json_helpers import writefile
 

modules/sd_hijack_freeu.py

@@ -156,15 +156,15 @@ def apply_freeu(p):
     s1 = p.freeu_s1 if p.freeu_s1 is not None else shared.opts.freeu_s1
     s2 = p.freeu_s2 if p.freeu_s2 is not None else shared.opts.freeu_s2
     state_freeu_b1, state_freeu_b2, state_freeu_s1, state_freeu_s2 = b1, b2, s1, s2
-    if hasattr(p.sd_model, 'enable_freeu'):
+    if hasattr(shared.sd_model, 'enable_freeu'):
         if enabled:
             freeu_device = get_fft_device()
             if freeu_device != devices.cpu:
                 p.extra_generation_params['FreeU'] = f'b1={b1} b2={b2} s1={s1} s2={s2}'
-                p.sd_model.enable_freeu(s1=s1, s2=s2, b1=b1, b2=b2)
+                shared.sd_model.enable_freeu(s1=s1, s2=s2, b1=b1, b2=b2)
                 state_enabled = True
         elif state_enabled:
-            p.sd_model.disable_freeu()
+            shared.sd_model.disable_freeu()
             state_enabled = False
     if enabled and state_enabled:
         log.info(f'Applying Free-U: b1={b1} b2={b2} s1={s1} s2={s2}')

modules/sd_hijack_hypertile.py

@@ -184,7 +184,7 @@ def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=128
 
 def context_hypertile_vae(p):
     from modules import shared
-    if p.sd_model is None or not get_opt(p, 'hypertile_vae_enabled'):
+    if shared.sd_model is None or not get_opt(p, 'hypertile_vae_enabled'):
         return nullcontext()
     if shared.opts.cross_attention_optimization == 'Sub-quadratic':
         log.warning('Hypertile UNet is not compatible with Sub-quadratic cross-attention optimization')
@@ -194,7 +194,7 @@ def context_hypertile_vae(p):
     error_reported = False
     set_resolution(p)
     max_h, max_w = 0, 0
-    vae = getattr(p.sd_model, "vae", None)
+    vae = getattr(shared.sd_model, "vae", None)
     if height == 0 or width == 0:
         log.warning('Hypertile VAE disabled: resolution unknown')
         return nullcontext()
@@ -215,7 +215,7 @@ def context_hypertile_vae(p):
 
 def context_hypertile_unet(p):
     from modules import shared
-    if p.sd_model is None or not get_opt(p, 'hypertile_unet_enabled'):
+    if shared.sd_model is None or not get_opt(p, 'hypertile_unet_enabled'):
         return nullcontext()
     if shared.opts.cross_attention_optimization == 'Sub-quadratic' and not shared.cmd_opts.experimental:
         log.warning('Hypertile UNet is not compatible with Sub-quadratic cross-attention optimization')
@@ -224,7 +224,7 @@ def context_hypertile_unet(p):
     error_reported = False
     set_resolution(p)
     max_h, max_w = 0, 0
-    unet = getattr(p.sd_model, "unet", None)
+    unet = getattr(shared.sd_model, "unet", None)
     if height == 0 or width == 0:
         log.warning('Hypertile VAE disabled: resolution unknown')
         return nullcontext()

modules/sd_models.py

@@ -18,7 +18,6 @@ from modules.modeldata import model_data
 from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoint_titles, get_closest_checkpoint_match, update_model_hashes, write_metadata, checkpoints_list # pylint: disable=unused-import
 from modules.sd_offload import get_module_names, disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import
 from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, apply_function_to_model, read_state_dict, get_state_dict_from_checkpoint # pylint: disable=unused-import
-from modules.secrets_manager import sanitize_dict, SecretStr
 
 
 model_dir = "Stable-diffusion"
@@ -314,7 +313,7 @@ def load_diffuser_initial(diffusers_load_config, op='model'):
         log.info(f'Load model {op}: path="{model_name}"')
         model_file = modelloader.download_diffusers_model(hub_id=model_name, variant=diffusers_load_config.get('variant', None))
         try:
-            log.debug(f'Load {op}: config={sanitize_dict(diffusers_load_config)}')
+            log.debug(f'Load {op}: config={diffusers_load_config}')
             sd_model = diffusers.DiffusionPipeline.from_pretrained(model_file, **diffusers_load_config)
         except Exception as e:
             log.error(f'Failed loading model: {model_file} {e}')
@@ -640,9 +639,9 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con
             diffusers_load_config.pop('requires_safety_checker', None)
             diffusers_load_config.pop('config_files', None)
             diffusers_load_config.pop('local_files_only', None)
-            log.debug(f'Setting {op}: pipeline={sd_model.__class__.__name__} config={sanitize_dict(diffusers_load_config)}') # pylint: disable=protected-access
+            log.debug(f'Setting {op}: pipeline={sd_model.__class__.__name__} config={diffusers_load_config}') # pylint: disable=protected-access
     except Exception as e:
-        log.error(f'Load {op}: file="{checkpoint_info.path}" pipeline={shared.opts.diffusers_pipeline} config={sanitize_dict(diffusers_load_config)} {e}')
+        log.error(f'Load {op}: file="{checkpoint_info.path}" pipeline={shared.opts.diffusers_pipeline} config={diffusers_load_config} {e}')
         if 'Weights for this component appear to be missing in the checkpoint' in str(e):
             log.error(f'Load {op}: file="{checkpoint_info.path}" is not a complete model')
         else:
@@ -799,9 +798,8 @@ def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: di
         "requires_safety_checker": False, # sd15 specific but we cant know ahead of time
         # "use_safetensors": True,
     }
-    hf_token = shared.opts.huggingface_token
-    if hf_token:
-        diffusers_load_config['token'] = SecretStr(hf_token)
+    if shared.opts.huggingface_token and len(shared.opts.huggingface_token) > 0:
+        diffusers_load_config['token'] = shared.opts.huggingface_token
     if revision is not None:
         diffusers_load_config['revision'] = revision
     if shared.opts.diffusers_model_load_variant != 'default':
@@ -1398,33 +1396,6 @@ def unload_model_weights(op='model'):
         log.debug(f'Unload {op}: {memory_stats()}  fn={fn}')
 
 
-def unload_auxiliary_models():
-    unloaded = []
-    try:
-        from modules.caption import vqa, openclip, tagger
-        inst = vqa.get_instance()
-        if inst.model is not None:
-            inst.unload()
-            unloaded.append('vqa')
-        openclip.unload_clip_model()
-        tagger.unload_model()
-    except Exception:
-        pass
-    try:
-        from modules.scripts_manager import scripts_txt2img
-        if scripts_txt2img is not None:
-            instances = [s for s in scripts_txt2img.scripts if 'prompt_enhance.py' in s.filename]
-            for inst in instances:
-                if inst.llm is not None:
-                    inst.unload()
-                    unloaded.append('prompt-enhance')
-    except Exception:
-        pass
-    if unloaded:
-        devices.torch_gc(force=True, reason='unload auxiliary')
-        log.debug(f'Unload auxiliary: {", ".join(unloaded)}')
-
-
 def hf_auth_check(checkpoint_info, force:bool=False):
     if shared.opts.offline_mode:
         log.info('Offline mode: skipping auth check')

scripts/prompt_enhance.py

@@ -260,6 +260,15 @@ class Options:
         'coder3101/Ministral-3-8B-Reasoning-2512-heretic': {},
         'coder3101/Ministral-3-14B-Reasoning-2512-heretic': {},
     }
+    models_cls = {
+        'qwen3_5': 'Qwen3_5ForConditionalGeneration',
+        'qwen3_5_moe': 'Qwen3_5MoeForConditionalGeneration',
+        'qwen3_vl': 'Qwen3VLForConditionalGeneration',
+        'qwen2_5_vl': 'Qwen2_5_VLForConditionalGeneration',
+        'qwen2_vl': 'Qwen2VLForConditionalGeneration',
+        'mistral3': 'Mistral3ForConditionalGeneration',
+    }
+
     # default = list(models)[1] # gemma-3-4b-it
     default = 'google/gemma-3-4b-it'
     supported = list(transformers.integrations.ggml.GGUF_CONFIG_MAPPING)
@@ -363,15 +372,12 @@ class Script(scripts_manager.Script):
 
             model_config = transformers.AutoConfig.from_pretrained(load_args['pretrained_model_name_or_path'], trust_remote_code=True, cache_dir=shared.opts.hfcache_dir)
             model_type = getattr(model_config, 'model_type', '')
-            model_type_cls = {
-                'qwen3_5': transformers.Qwen3_5ForConditionalGeneration,
-                'qwen3_5_moe': transformers.Qwen3_5MoeForConditionalGeneration,
-                'qwen3_vl': transformers.Qwen3VLForConditionalGeneration,
-                'qwen2_5_vl': transformers.Qwen2_5_VLForConditionalGeneration,
-                'qwen2_vl': transformers.Qwen2VLForConditionalGeneration,
-                'mistral3': transformers.Mistral3ForConditionalGeneration,
-            }
-            cls_name = model_type_cls.get(model_type, transformers.AutoModelForCausalLM)
+            cls_name = transformers.AutoModelForCausalLM
+            custom_cls_name = self.options.models_cls.get(model_type, None)
+            if custom_cls_name:
+                custom_cls = getattr(transformers, custom_cls_name, None)
+                if custom_cls:
+                    cls_name = custom_cls
 
             sd_models.set_caption_load_options()
             try:
@@ -396,10 +402,7 @@ class Script(scripts_manager.Script):
             tokenizer_args = { 'pretrained_model_name_or_path': model_repo }
             if model_tokenizer:
                 tokenizer_args['subfolder'] = model_tokenizer
-            self.tokenizer = cls.from_pretrained(
-                **tokenizer_args,
-                cache_dir=shared.opts.hfcache_dir,
-            )
+            self.tokenizer = cls.from_pretrained(**tokenizer_args, cache_dir=shared.opts.hfcache_dir)
             self.tokenizer.is_processor = model_repo in self.options.img2img
 
             if debug_enabled: