Intergrate fp16/bf16 support to sdxl model loading

[Isotr0py]

Maybe related issue: #788

• Specify fp16/bf16 dtype when loading model state_dict
• Directly load unet and vae to fp16/bf16 dtype if use full fp16/bf16.
  -text_encoders are remained to fp32 for cache outputs.

This should reduce RAM/VRAM usage peak when enable --full_fp16/--full_bf16 during model loading on CPU/GPU.

It also reduced RAM usage when loading checkpoint from safetensors format.

[kohya-ss]

Thank you for this! This is very useful. I wonder it is possible to load Text Encoders as fp32, even with full_fp16/bf16 option. If it is not possible, this feature may be enabled when lowram/lowvram option is specified.

[Isotr0py]

Yes, dtype is not passed to _load_state_dict_on_device for text_encoders. So text_encoders will always load as fp32, no matter whether full_fp16/bf16 is enabled or not.

[kohya-ss]

Thank you for clarification, and sorry for misunderstood. I understood the text encoders are loaded in fp16/bf16 only if the Diffusers format is used.

[kohya-ss]

I've merged. This significantly reduces the peak memory usage. Thanks again!

[FurkanGozukara]

could this have introduced a major bug?

My SDXL LoRAs are super trained now with same settings compared to before this change

Actually I am testing right now effect of this option