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']: # Added 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]