what function does forge use to load models

[Enferlain]

I wanted to add support for an a1111 merger extension, but I can't figure out the model loading part

on a1111 it was just

sd_models.load_model(sd_models.CheckpointInfo(model_path))

in forge I tried /sdapi/v1/options {sd_model_checkpoint: modelname.safetensors} which for some reason says it can't find the model even though it's there at the time the error appears

File "D:\stable-diffusion-webui-forge\modules\api\api.py", line 677, in set_config
    raise RuntimeError(f"model {checkpoint_name!r} not found")
RuntimeError: model 'D:\stable-diffusion-webui-forge\models\Stable-diffusion\bbwm-2-mymixGJem_xlsQromew-thisCameToMeInADream_mid-it_1.safetensors' not found

and currently I was trying with a custom function in the extension like this

def load_model_forge(model_path: str):
    """Custom model loading function for Forge."""

    try:
        # 1. Unload existing model
        memory_management.unload_all_models()
        memory_management.soft_empty_cache()
        gc.collect()

        # 2. Prepare model data
        checkpoint_info = sd_models.CheckpointInfo(model_path)
        additional_state_dicts = []  # Modify as needed based on your requirements

        # 3. Load the model
        sd_model = forge_loader(checkpoint_info.filename, additional_state_dicts=additional_state_dicts)

        # 4. Initialize model attributes
        sd_model.extra_generation_params = {}
        sd_model.comments = []
        sd_model.sd_checkpoint_info = checkpoint_info
        sd_model.filename = checkpoint_info.filename
        sd_model.sd_model_hash = checkpoint_info.calculate_shorthash()

        # 5. Trigger callbacks
        script_callbacks.model_loaded_callback(sd_model)

        # Update the shared.sd_model object in Forge
        shared.sd_model = sd_model

        logger.info(f"Model loaded successfully from: {model_path} in Forge")

    except Exception as e:
        logger.error(f"Error loading model in Forge: {e}", exc_info=True)
        raise

Which does load the model, but the subsequent generation request makes images that are burned in greyscale and have no discernible features. The same model will make normal images when manually generating in the ui, and same for when just sending the generate request, so I assume it's related to the loading part.

if someone knows pls help

[lllyasviel]

current logic is modules.sd_model.SdModelData.forge_loading_parameters stores cached loading parameters
you can globally search "forge_loading_parameters" to find all codes that write it
when t2i/i2i process starts, it compare if the loading_parameters is different from the current loaded model, if different, then load with sd_models.forge_model_reload

so if you do not need to have the model instance immediately, you can just write "forge_loading_parameters" . The model will be load after the next time user click generate. This is highly recommended since randomly load models may make user experience very bad in the Flux era when loading a model can take 40 seconds.

If you must have model instance immediately, you can just write "forge_loading_parameters" and then call sd_models.forge_model_reload immediately after that

[Enferlain]

Ok I replaced the custom approach with directly calling the functions you recommended and it worked for the loading part

        elif webui == "forge":
            # Construct the checkpoint_info dictionary
            checkpoint_info = sd_models.CheckpointInfo(model_path)

            # Update forge_loading_parameters
            sd_models.model_data.forge_loading_parameters = {
                "checkpoint_info": checkpoint_info}  # Add other parameters as needed

            # Immediately load the model
            sd_models.forge_model_reload()
            print(f"Bayesian Merger: Loaded model from {model_path} in Forge")

But I'm unfortunately encountering the thing I mentioned with generate requests coming through the api looking like this

yet same model does fine when clicking generate. I tried with both forge_model_reload and without.

The process would go like: merge model -> send model to forge -> gen images through generate request -> evaluate -> repeat

The steps individually seem to work fine, but combining sending the model and then the generate request seems to not play well for some reason. Since the generate request works on models that were selected from the checkpoint list (or maybe if I just manually press the generate button?) I'm guessing something else would need to be included in the way I'm trying to load models.

[Enferlain]

Actually I think this works. Adding refresh_model_loading_parameters seems to have fixed it so far.

            # Directly update forge_loading_parameters
            sd_models.model_data.forge_loading_parameters = {
                "checkpoint_info": sd_models.CheckpointInfo(model_path),  # Create a CheckpointInfo object
                "additional_modules": []  # Add other parameters as needed
            }

            # Call forge_model_reload to load the model
            sd_models.forge_model_reload()

            # Call refresh_model_loading_parameters to update settings
            refresh_model_loading_parameters()

Thanks for the clues

edit ^ this was almost right, had to use checkpoint_change instead of just raw refresh_model_loading_parameters otherwise even if the correct model is loaded, it would be overwritten by the "selected model" at gen request.

            # Set the loaded model as the active model using checkpoint_change
            main_entry.checkpoint_change(os.path.basename(model_path))

gonna leave this here just in case someone needs it in the future

[Enferlain]

@lllyasviel what about for unloading? I tried free_memory and unload_all_models and the 1st one does print some unload related messages about trying to free memory in the console, but it doesn't unload the model, and the 2nd one doesn't seem to have an effect. Asking cuz in the api part it still uses sd_models.unload_model_weights which is just a pass in forge.