UNet Extractor and Remover for Stable Diffusion

[ali0une]

Hi there.

it could be interesting to have the code in Forge in order to use these types of checkpoints that are shipped "UNET Only" and are way lighter.

Related to this there is a project on github to extract UNET from checkpoints :
https://github.com/captainzero93/extract-unet-safetensor

Just imagine the space freed on our hard drives ...

[ali0une]

OK ... so there was a bug with detection for flux unet in the code of extract-unet-safetensor

So i gave it another try and it works in Forge!

(venv) user@box:~/whatever/extract-unet-safetensor$ python UNetExtractor.py /home/user/whatever/stable-diffusion-webui/models/Stable-diffusion/flux/flux1-dev-fp8.safetensors ./outputs/flux1-dev-fp8_unet.safetensors ./outputs/flux1-dev-fp8_no_unet.safetensors --model_type flux
2024-08-17 21:39:35,613 - INFO - PyTorch version: 2.4.0+cu121
2024-08-17 21:39:35,613 - INFO - CUDA version: 12.1
2024-08-17 21:39:35,632 - INFO - GPU device name: NVIDIA GeForce RTX 3090
2024-08-17 21:39:35,632 - INFO - Processing /home/user/whatever/stable-diffusion-webui/models/Stable-diffusion/flux/flux1-dev-fp8.safetensors on cuda
2024-08-17 21:39:35,632 - INFO - Model type: flux
2024-08-17 21:39:37,709 - INFO - Total tensors processed: 1442
2024-08-17 21:39:37,709 - INFO - UNet tensors: 780
2024-08-17 21:39:37,709 - INFO - Non-UNet tensors: 662
2024-08-17 21:39:37,709 - INFO - Unique key prefixes found: model, text_encoders, vae
2024-08-17 21:39:37,958 - INFO - Saving extracted UNet to outputs/flux1-dev-fp8_unet.safetensors
2024-08-17 21:39:44,178 - INFO - Saving model without UNet to outputs/flux1-dev-fp8_no_unet.safetensors
2024-08-17 21:39:47,147 - INFO - Processing complete!
(venv) user@box:~/whatever/extract-unet-safetensor$ ls -la outputs/
total 17G
drwxr-xr-x 2 user user 4,0K 2024-08-17 21:47:36 .
drwxr-xr-x 5 user user 4,0K 2024-08-17 21:37:25 ..
-rw-r--r-- 1 user user 5,0G 2024-08-17 21:39:47 flux1-dev-fp8_no_unet.safetensors
-rw-r--r-- 1 user user  12G 2024-08-17 21:39:44 flux1-dev-fp8_unet.safetensors

You end up with a 12G flux1-dev-fp8_unet.safetensors (against 17.2G of the original flux1-dev-fp8.safetensors) so it's 5Go freed ... imagine doing this with all our SD/SDXL checkpoints!

Then you put this flux1-dev-fp8_unet.safetensors in your Stable-diffusion/models folder and you select it for inference in Forge plus you select VAE and Text Encoders (VAE : flux1-dev-vae-float16.safetensors / Text Encoders : flux1-dev-clip_l.safetensors + flux1-dev-t5xxl_fp16.safetensors)

Enter a prompt, hit generate ... profit!

Not sure about how to make it work with SDXL models, tried before and ended up with errors ... will check later.