| import gradio as gr
|
| from modules import scripts, shared, sd_models, lowvram, devices, paths
|
| import gc
|
| import torch
|
| import os
|
|
|
| try:
|
| from modules.sd_models import forge_model_reload, model_data, CheckpointInfo
|
| from modules_forge.main_entry import forge_unet_storage_dtype_options
|
| from backend.memory_management import free_memory as forge_free_memory
|
| from modules.timer import Timer
|
| forge = True
|
| except ImportError:
|
| forge = False
|
| class CheckpointInfo:
|
| def __init__(self, filename):
|
| self.filename = filename
|
| self.name = os.path.splitext(os.path.basename(filename))[0]
|
| self.name_or_path = filename
|
| self.sha256 = None
|
| self.ids = None
|
| self.model_name = self.name
|
| self.title = self.name
|
| class Timer:
|
| def record(self, *args, **kwargs): pass
|
|
|
| class ModelUtilState:
|
| last_loaded_checkpoint_info_dict = None
|
| last_forge_model_params = None
|
| is_model_unloaded_by_ext = False
|
|
|
| state = ModelUtilState()
|
|
|
| def get_current_checkpoint_info():
|
| if forge and hasattr(model_data, 'sd_checkpoint_info') and model_data.sd_checkpoint_info:
|
| return model_data.sd_checkpoint_info
|
| if hasattr(shared, 'sd_model') and shared.sd_model and hasattr(shared.sd_model, 'sd_checkpoint_info') and shared.sd_model.sd_checkpoint_info:
|
| return shared.sd_model.sd_checkpoint_info
|
| if shared.opts.sd_model_checkpoint:
|
| checkpoint_path = sd_models.get_checkpoint_path(shared.opts.sd_model_checkpoint)
|
| if checkpoint_path:
|
| return CheckpointInfo(checkpoint_path)
|
| return None
|
|
|
| def checkpoint_info_to_dict(chkpt_info):
|
| if not chkpt_info:
|
| return None
|
| return {
|
| "filename": getattr(chkpt_info, 'filename', None),
|
| "name": getattr(chkpt_info, 'name', None),
|
| "name_or_path": getattr(chkpt_info, 'name_or_path', getattr(chkpt_info, 'filename', None)),
|
| "sha256": getattr(chkpt_info, 'sha256', None),
|
| "model_name": getattr(chkpt_info, 'model_name', None),
|
| "title": getattr(chkpt_info, 'title', None),
|
| }
|
|
|
| def ensure_name_or_path(info_obj):
|
| if not info_obj:
|
| return info_obj
|
| if not hasattr(info_obj, 'name_or_path') or not getattr(info_obj, 'name_or_path', None):
|
| filename_attr = getattr(info_obj, 'filename', None)
|
| title_attr = getattr(info_obj, 'title', None)
|
| name_attr = getattr(info_obj, 'name', None)
|
| if filename_attr:
|
| print(f"Info object was missing 'name_or_path'. Setting from 'filename': {filename_attr}")
|
| info_obj.name_or_path = filename_attr
|
| elif title_attr:
|
| print(f"Info object was missing 'name_or_path/filename'. Setting from 'title': {title_attr}")
|
| info_obj.name_or_path = title_attr
|
| elif name_attr:
|
| print(f"Info object was missing 'name_or_path/filename/title'. Setting from 'name': {name_attr}")
|
| info_obj.name_or_path = name_attr
|
| else:
|
| print(f"CRITICAL: Info object is missing 'name_or_path', 'filename', 'title', and 'name'. Cannot reliably set 'name_or_path'.")
|
| return info_obj
|
|
|
| def dict_to_checkpoint_info(chkpt_dict):
|
| if not chkpt_dict or not chkpt_dict.get('name_or_path'):
|
| print(f"Warning: chkpt_dict is invalid or missing 'name_or_path': {chkpt_dict}")
|
| return None
|
|
|
| target_model_identifier = chkpt_dict['name_or_path']
|
| print(f"Attempting to find CheckpointInfo for: {target_model_identifier}")
|
|
|
| available_checkpoints = sd_models.checkpoints_list
|
| found_info = None
|
|
|
| for name, info_obj_from_list in available_checkpoints.items():
|
| info_name_or_path = getattr(info_obj_from_list, 'name_or_path', None)
|
| info_filename = getattr(info_obj_from_list, 'filename', None)
|
| info_title = getattr(info_obj_from_list, 'title', None)
|
| match_found = False
|
| if info_name_or_path and info_name_or_path == target_model_identifier: match_found = True
|
| elif info_filename and info_filename == target_model_identifier: match_found = True
|
| elif name == target_model_identifier: match_found = True
|
| elif info_title and info_title == target_model_identifier: match_found = True
|
| if match_found:
|
| print(f"Found matching CheckpointInfo in available_checkpoints: {name}")
|
| found_info = info_obj_from_list
|
| break
|
| if found_info:
|
| return ensure_name_or_path(found_info)
|
|
|
| print(f"CheckpointInfo for '{target_model_identifier}' not found in list. Attempting to create new one.")
|
| if os.path.exists(target_model_identifier):
|
| print(f"File exists at path: {target_model_identifier}. Creating new CheckpointInfo.")
|
| newly_created_info = CheckpointInfo(target_model_identifier)
|
| for key, value in chkpt_dict.items():
|
| if not hasattr(newly_created_info, key) or getattr(newly_created_info, key) is None:
|
| setattr(newly_created_info, key, value)
|
| return ensure_name_or_path(newly_created_info)
|
| else:
|
| print(f"File does not exist at path: {target_model_identifier}. Cannot create CheckpointInfo.")
|
|
|
| print(f"Warning: Could not reconstruct CheckpointInfo for {target_model_identifier}.")
|
| return None
|
|
|
|
|
| def unload_model_logic():
|
| model_loaded = (forge and hasattr(model_data, 'sd_model') and model_data.sd_model) or \
|
| (not forge and hasattr(shared, 'sd_model') and shared.sd_model)
|
| if not model_loaded:
|
| state.is_model_unloaded_by_ext = False
|
| return "Model is already unloaded or not loaded."
|
|
|
| print("Unloading SD model...")
|
| current_info = get_current_checkpoint_info()
|
| if current_info:
|
| state.last_loaded_checkpoint_info_dict = checkpoint_info_to_dict(current_info)
|
| print(f"Storing info for model: {state.last_loaded_checkpoint_info_dict.get('name_or_path')}")
|
| else:
|
| state.last_loaded_checkpoint_info_dict = None
|
| print("Could not get current checkpoint info to store.")
|
|
|
| if forge:
|
| if hasattr(model_data, "forge_loading_parameters") and model_data.forge_loading_parameters:
|
| state.last_forge_model_params = model_data.forge_loading_parameters.copy()
|
| else:
|
| state.last_forge_model_params = None
|
| sd_models.model_data.sd_model = None
|
| if hasattr(sd_models.model_data, 'loaded_sd_models'):
|
| sd_models.model_data.loaded_sd_models = []
|
| if hasattr(sd_models.model_data, 'forge_objects'):
|
| for attr in ['unet', 'vae', 'clip_l', 'clip_g', 'clip_vision', 'gligen', 'controlnet_predict', 'patch_manager', 'conditioner']:
|
| if hasattr(sd_models.model_data.forge_objects, attr):
|
| setattr(sd_models.model_data.forge_objects, attr, None)
|
| cuda_device_str = devices.get_cuda_device_string() if torch.cuda.is_available() else "cpu"
|
| if torch.cuda.is_available():
|
| forge_free_memory(torch.cuda.memory_allocated(cuda_device_str), cuda_device_str, free_all=True)
|
| print("Forge model components cleared and memory freed.")
|
| else:
|
| sd_models.unload_model_weights()
|
| print("Standard model unloaded.")
|
|
|
| lowvram.module_in_gpu = None
|
| shared.sd_model = None
|
| gc.collect()
|
| if torch.cuda.is_available():
|
| torch.cuda.empty_cache()
|
| state.is_model_unloaded_by_ext = True
|
| return "Model unloaded successfully. VRAM freed."
|
|
|
| def _ensure_module_on_device(module, module_name, target_device, indent=" "):
|
| if module and isinstance(module, torch.nn.Module) and next(module.parameters(), None) is not None:
|
| current_device = next(module.parameters()).device
|
| if current_device.type != target_device.type or (target_device.type == 'cuda' and current_device.index != target_device.index):
|
| print(f"{indent}Moving {module_name} from {current_device} to {target_device}...")
|
| module.to(target_device)
|
| return True
|
| return False
|
|
|
| def reload_last_model_logic():
|
| model_currently_loaded = (forge and hasattr(model_data, 'sd_model') and model_data.sd_model and model_data.sd_model is not shared.sd_model_empty) or \
|
| (not forge and hasattr(shared, 'sd_model') and shared.sd_model and shared.sd_model is not shared.sd_model_empty)
|
|
|
| if model_currently_loaded and not state.is_model_unloaded_by_ext:
|
| return "Model is already loaded and was not unloaded by this extension. No action taken."
|
|
|
| if not state.last_loaded_checkpoint_info_dict:
|
| if shared.opts.sd_model_checkpoint:
|
| print(f"No specific model info stored by extension, trying to use WebUI's selected model: {shared.opts.sd_model_checkpoint}")
|
| checkpoint_path = sd_models.get_checkpoint_path(shared.opts.sd_model_checkpoint)
|
| if checkpoint_path:
|
| state.last_loaded_checkpoint_info_dict = checkpoint_info_to_dict(CheckpointInfo(checkpoint_path))
|
| else:
|
| return "No last model information found and WebUI's selected model could not be resolved."
|
| else:
|
| return "No last model information found to reload."
|
|
|
| chkpt_info_to_load = dict_to_checkpoint_info(state.last_loaded_checkpoint_info_dict)
|
| if not chkpt_info_to_load or not getattr(chkpt_info_to_load, 'name_or_path', None):
|
| return f"Could not reconstruct valid CheckpointInfo from stored data: {state.last_loaded_checkpoint_info_dict}. Cannot reload."
|
|
|
| model_display_name = getattr(chkpt_info_to_load, 'name_or_path', getattr(chkpt_info_to_load, 'filename', 'Unknown Model'))
|
| print(f"Reloading SD model: {model_display_name}")
|
|
|
| try:
|
| devices.torch_gc()
|
|
|
| if forge:
|
| print("Forge: Reloading using forge_model_reload()...")
|
| if state.last_forge_model_params:
|
| sd_models.model_data.forge_loading_parameters = state.last_forge_model_params.copy()
|
| sd_models.model_data.forge_loading_parameters['checkpoint_info'] = chkpt_info_to_load
|
| else:
|
| print("Warning: No specific Forge params stored, building defaults for reload.")
|
| unet_storage_dtype, _ = forge_unet_storage_dtype_options.get(shared.opts.forge_unet_storage_dtype, (None, False))
|
| sd_models.model_data.forge_loading_parameters = dict(
|
| checkpoint_info=chkpt_info_to_load,
|
| additional_modules=shared.opts.forge_additional_modules,
|
| unet_storage_dtype=unet_storage_dtype
|
| )
|
| sd_models.model_data.forge_hash = None
|
|
|
| forge_model_reload()
|
|
|
| if not sd_models.model_data.sd_model:
|
| raise RuntimeError("forge_model_reload() did not populate model_data.sd_model.")
|
|
|
| shared.sd_model = sd_models.model_data.sd_model
|
| print("Forge: forge_model_reload() completed.")
|
|
|
| if torch.cuda.is_available():
|
| cuda_device = torch.device(devices.get_cuda_device_string())
|
| print(f"Forge: Verifying device placement on {cuda_device} after reload...")
|
|
|
| _ensure_module_on_device(shared.sd_model, "shared.sd_model (main)", cuda_device)
|
|
|
| if hasattr(shared.sd_model, 'forge_objects') and shared.sd_model.forge_objects:
|
| fo = shared.sd_model.forge_objects
|
| _ensure_module_on_device(getattr(fo, 'unet', None), "UNet (from forge_objects)", cuda_device)
|
| _ensure_module_on_device(getattr(fo, 'vae', None), "VAE (from forge_objects)", cuda_device)
|
| _ensure_module_on_device(getattr(fo, 'clip', None), "CLIP (main from forge_objects)", cuda_device)
|
| if hasattr(fo, 'clip') and fo.clip:
|
| _ensure_module_on_device(getattr(fo.clip,'cond_stage_model', None), "CLIP cond_stage_model", cuda_device)
|
|
|
| if hasattr(shared.sd_model, 'conditioner') and shared.sd_model.conditioner:
|
| _ensure_module_on_device(shared.sd_model.conditioner, "Conditioner", cuda_device)
|
| if hasattr(shared.sd_model.conditioner, 'embedders'):
|
| for i, embedder in enumerate(shared.sd_model.conditioner.embedders):
|
| _ensure_module_on_device(embedder, f"Embedder {i}", cuda_device)
|
|
|
| print("Forge: Device verification and correction attempt finished.")
|
| else:
|
| sd_models.load_model(chkpt_info_to_load)
|
| print("Standard model reloaded.")
|
| if torch.cuda.is_available() and shared.sd_model:
|
| cuda_device = torch.device(devices.get_cuda_device_string())
|
| _ensure_module_on_device(shared.sd_model, "shared.sd_model (main)", cuda_device)
|
|
|
|
|
| state.is_model_unloaded_by_ext = False
|
| return f"Model '{model_display_name}' reloaded successfully."
|
|
|
| except Exception as e:
|
| print(f"Error reloading model: {e}")
|
| import traceback
|
| traceback.print_exc()
|
| lowvram.module_in_gpu = None
|
| shared.sd_model = None
|
| if forge and hasattr(model_data, 'sd_model'): model_data.sd_model = None
|
| gc.collect()
|
| if torch.cuda.is_available(): torch.cuda.empty_cache()
|
| return f"Error reloading model: {e}. Model remains unloaded."
|
|
|
|
|
| class UnloadReloadModelScript(scripts.Script):
|
| def title(self):
|
| return "Model Unload/Reload Util"
|
|
|
| def show(self, is_img2img):
|
| return scripts.AlwaysVisible
|
|
|
| def ui(self, is_img2img):
|
| with gr.Accordion(self.title(), open=False):
|
| with gr.Row():
|
| unload_button = gr.Button("Unload Current SD Model (Free VRAM)")
|
| reload_button = gr.Button("Reload Last Unloaded SD Model")
|
| status_text = gr.Textbox(label="Status", value="Ready.", interactive=False, lines=3, max_lines=3)
|
|
|
| unload_button.click(fn=unload_model_logic, inputs=[], outputs=[status_text])
|
| reload_button.click(fn=reload_last_model_logic, inputs=[], outputs=[status_text])
|
| return [unload_button, reload_button, status_text] |
