ComfyUI Node: Model Pruner (mtb)

Model Pruner (mtb):

The Model Pruner (mtb) node is designed to optimize and streamline your AI models by pruning unnecessary components and converting model precisions. This node helps in reducing the model size and improving performance by removing redundant parts and adjusting the precision of tensors. It supports operations on different parts of the model, including UNet, CLIP, and VAE, and can handle various precision formats such as FP8, FP16, BF16, and FP32. The main goal of this node is to enhance the efficiency of your models, making them faster and more resource-efficient without compromising their performance.

Model Pruner (mtb) Input Parameters:

save_separately

This parameter determines whether the pruned model components should be saved separately. If set to True, each part of the model (UNet, CLIP, VAE) will be saved in its own file. This can be useful for modularity and easier management of model components. The default value is False.

save_folder

Specifies the directory where the pruned model files will be saved. This should be a valid path on your filesystem where you have write permissions. The default value is an empty string, which means the current working directory will be used.

fix_clip

A boolean parameter that, when set to True, applies fixes to the CLIP component of the model. This is useful for correcting known issues with CLIP models. The default value is False.

remove_junk

This parameter indicates whether to remove junk data from the model. Junk data can include unnecessary keys or values that do not contribute to the model's performance. Setting this to True helps in further reducing the model size. The default value is False.

ema_mode

Specifies the mode for handling Exponential Moving Average (EMA) in the model. This can be set to different modes depending on how you want to manage EMA weights. The default value is an empty string.

precision_unet

Defines the precision format for the UNet component. Supported values include FP8, FP16, BF16, and FP32. This parameter helps in converting the UNet tensors to the specified precision, optimizing memory usage and computational efficiency.

precision_clip

Defines the precision format for the CLIP component. Supported values include FP8, FP16, BF16, and FP32. This parameter helps in converting the CLIP tensors to the specified precision, optimizing memory usage and computational efficiency.

precision_vae

Defines the precision format for the VAE component. Supported values include FP8, FP16, BF16, and FP32. This parameter helps in converting the VAE tensors to the specified precision, optimizing memory usage and computational efficiency.

operation_unet

Specifies the operation to be performed on the UNet component. Possible values are CONVERT, COPY, and DELETE. This parameter determines whether to convert the precision, copy as is, or delete the UNet component.

operation_clip

Specifies the operation to be performed on the CLIP component. Possible values are CONVERT, COPY, and DELETE. This parameter determines whether to convert the precision, copy as is, or delete the CLIP component.

operation_vae

Specifies the operation to be performed on the VAE component. Possible values are CONVERT, COPY, and DELETE. This parameter determines whether to convert the precision, copy as is, or delete the VAE component.

unet

A dictionary containing the UNet model tensors. This parameter is optional and can be None if the UNet component is not being pruned or modified.

clip

A dictionary containing the CLIP model tensors. This parameter is optional and can be None if the CLIP component is not being pruned or modified.

vae

A dictionary containing the VAE model tensors. This parameter is optional and can be None if the VAE component is not being pruned or modified.

Model Pruner (mtb) Output Parameters:

None

This node does not produce any direct output parameters. The results of the pruning and conversion operations are saved to the specified directory.

Model Pruner (mtb) Usage Tips:

• Ensure that the save_folder parameter is set to a valid directory where you have write permissions to avoid any file saving errors.
• Use the fix_clip parameter if you are aware of specific issues with your CLIP model that need correction.
• Adjust the precision_unet, precision_clip, and precision_vae parameters to optimize the model's memory usage and performance based on your hardware capabilities.
• Utilize the remove_junk parameter to further reduce the model size by eliminating unnecessary data.

Model Pruner (mtb) Common Errors and Solutions:

Cannot convert {tensor.dtype} to fp8

• Explanation: This error occurs when the tensor's data type is not compatible with FP8 precision.
• Solution: Ensure that the tensor data types are compatible with the specified precision format or choose a different precision format.

Not a torch tensor, skipping key

• Explanation: This error indicates that a non-tensor value was encountered in the model dictionary.
• Solution: Verify that all values in the model dictionaries (UNet, CLIP, VAE) are torch tensors.

[Converter] Fixed novelai error key {k}

• Explanation: This log message indicates that a known issue with a NovelAI key was detected and fixed.
• Solution: No action needed; this is an informational message confirming that the issue was automatically corrected.