captainzero93 / extract-unet-safetensor

Processes SafeTensors files for Stable Diffusion 1.5 (SD 1.5), Stable Diffusion XL (SDXL), and FLUX models. It extracts the UNet into a separate file and creates a new file with the remaining model components (without the UNet).
MIT License
43 stars 4 forks source link
diffusion flux python sdxl stable t5xxl unet

UNet Extractor and Remover for Stable Diffusion 1.5, SDXL, and FLUX

This Python script (UNetExtractor.py) processes SafeTensors files for Stable Diffusion 1.5 (SD 1.5), Stable Diffusion XL (SDXL), and FLUX models. It extracts the UNet into a separate file and creates a new file with the remaining model components (without the UNet).

FLUX Example

Above example: UNetExtractor.py flux1-dev.safetensors flux1-dev_unet.safetensors flux1-dev_non_unet.safetensors --model_type flux --verbose

AUTOMATIC1111 Extension for UNet Loading

We've developed an extension for AUTOMATIC1111's Stable Diffusion Web UI that allows you to load and use the extracted UNet files directly within the interface. This extension seamlessly integrates with the txt2img workflow, enabling you to utilize the space-saving benefits of separated UNet files without compromising on functionality.

Extension Features:

To use the extension, please visit our UNet Loader Extension Repository for installation and usage instructions.

Why UNet Extraction?

Using UNets instead of full checkpoints can save a significant amount of disk space, especially for models that utilize large text encoders. This is particularly beneficial for models like FLUX, which has a large number of parameters. Here's why:

This tool helps you extract UNets from full checkpoints, allowing you to take advantage of these space-saving benefits across SD 1.5, SDXL, and open-source FLUX models.

Features

Requirements

Installation

  1. Clone this repository or download the UNetExtractor.py script.

  2. It's recommended to create a new virtual environment:

    python -m venv unet_extractor_env
  3. Activate the virtual environment:

    • On Windows:
      unet_extractor_env\Scripts\activate
    • On macOS and Linux:
      source unet_extractor_env/bin/activate
  4. Install the required libraries with specific versions for debugging:

    pip install numpy==1.23.5 torch==2.0.1 safetensors==0.3.1
  5. If you're using CUDA, install the CUDA-enabled version of PyTorch:

    pip install torch==2.0.1+cu117 -f https://download.pytorch.org/whl/cu117/torch_stable.html

    Replace cu117 with your CUDA version (e.g., cu116, cu118) if different.

  6. Optionally, install psutil for enhanced system resource reporting:

    pip install psutil==5.9.0

Note: The versions above are examples and may need to be adjusted based on your system requirements and CUDA version. These specific versions are recommended for debugging purposes as they are known to work together. For regular use, you may use the latest versions of these libraries.

Usage

Run the script from the command line with the following syntax:

python UNetExtractor.py <input_file> <unet_output_file> <non_unet_output_file> --model_type <sd15|sdxl|flux> [--verbose] [--num_threads <num>] [--gpu_limit <percent>] [--cpu_limit <percent>]

Arguments

Examples

For Stable Diffusion 1.5 using CUDA (if available):

python UNetExtractor.py path/to/sd15_model.safetensors path/to/output_sd15_unet.safetensors path/to/output_sd15_non_unet.safetensors --model_type sd15 --verbose

For Stable Diffusion XL using CUDA (if available):

python UNetExtractor.py path/to/sdxl_model.safetensors path/to/output_sdxl_unet.safetensors path/to/output_sdxl_non_unet.safetensors --model_type sdxl --verbose

For FLUX models using CUDA (if available) with 8 threads and 80% GPU usage limit:

python UNetExtractor.py path/to/flux_model.safetensors path/to/output_flux_unet.safetensors path/to/output_flux_non_unet.safetensors --model_type flux --verbose --num_threads 8 --gpu_limit 80

How It Works

  1. The script checks for CUDA availability (if PyTorch is installed) and prompts to choose between CPU-only and GPU-assisted processing.
  2. It determines the optimal number of threads to use based on the system's CPU cores (if not manually specified).
  3. It opens the input SafeTensors file using the safetensors library.
  4. The script iterates through all tensors in the file, separating UNet-related tensors from other tensors.
  5. For SD 1.5 and FLUX models, it removes the "model.diffusion_model." prefix from UNet tensor keys.
  6. For SDXL, it keeps the original key names for both UNet and non-UNet tensors.
  7. The script uses multi-threading to process tensors concurrently, improving performance.
  8. GPU and CPU usage are limited based on user-specified percentages or default values.
  9. The extracted UNet tensors are saved to a new SafeTensors file.
  10. The remaining non-UNet tensors are saved to a separate SafeTensors file.
  11. RAM offloading is implemented to manage memory usage, especially for large models.

Using Extracted UNets with AUTOMATIC1111

After extracting UNet files using this tool, you can easily use them in AUTOMATIC1111's Stable Diffusion Web UI:

  1. Install our UNet Loader extension in your AUTOMATIC1111 setup.
  2. Place the extracted UNet and non-UNet files in the extension's designated folder.
  3. Use the extension's interface to select and load your desired UNet and non-UNet components.
  4. Generate images using txt2img as usual, now benefiting from the space-saving and flexibility of separated UNet files.

For detailed instructions, please refer to the UNet Loader Extension Repository.

Debugging Information

When running the script with the --verbose flag, you'll see detailed debugging information, including:

Example debug output:

2024-08-17 21:06:30,500 - DEBUG - Current UNet count: 770
2024-08-17 21:06:30,500 - DEBUG - ---
2024-08-17 21:06:31,142 - DEBUG - Processing key: vector_in.out_layer.weight
2024-08-17 21:06:31,142 - DEBUG - Tensor shape: torch.Size([3072, 3072])
2024-08-17 21:06:31,172 - DEBUG - Classified as non-UNet tensor
2024-08-17 21:06:31,172 - DEBUG - Current UNet count: 770
2024-08-17 21:06:31,172 - DEBUG - ---
2024-08-17 21:06:31,203 - INFO - Total tensors processed: 780
2024-08-17 21:06:31,203 - INFO - UNet tensors: 770
2024-08-17 21:06:31,203 - INFO - Non-UNet tensors: 10
2024-08-17 21:06:31,203 - INFO - Unique key prefixes found: double_blocks, final_layer, guidance_in, img_in, single_blocks, time_in, txt_in, vector_in

This output helps identify issues with tensor classification, resource usage, and overall processing flow.

Notes

Troubleshooting

If you encounter any issues:

  1. Ensure you're using the recommended library versions as specified in the Installation section.
  2. Run the script with the --verbose flag to get detailed debugging information.
  3. Check for compatibility between your CUDA version and the installed PyTorch version.
  4. If you encounter a NumPy version compatibility error with PyTorch, such as:
    A module that was compiled using NumPy 1.x cannot be run in
    NumPy 2.0.1 as it may crash.

    Try downgrading NumPy to version 1.23.5 as recommended in the installation instructions.

  5. Ensure you have the latest version of the safetensors library installed.
  6. Check that your input file is a valid SafeTensors file for the specified model type.
  7. Make sure you have read permissions for the input file and write permissions for the output directory.
  8. If you're having issues with CUDA, try running with CPU-only processing to see if it resolves the problem.
  9. If you encounter any "module not found" errors, ensure all required libraries are installed in your virtual environment.
  10. Check the console output for any error messages or stack traces that can help identify the issue.

If you continue to experience issues after trying these steps, please open an issue on the GitHub repository with details about your system configuration, the command you're using, and the full error message or debugging output.

Contributing

Contributions, issues, and feature requests are welcome! Feel free to check issues page if you want to contribute.

License

This project is licensed under the MIT License. See the LICENSE file for details.

Citation

If you use UNet Extractor and Remover in your research or projects, please cite it as follows:

[For commercial licensing cyberjunk77@gmail.com] (captainzero93). (2024). UNet Extractor and Remover for Stable Diffusion 1.5, SDXL, and FLUX. GitHub. https://github.com/captainzero93/unet-extractor

Acknowledgements