Repository navigation
override Flux transformer - #639
Conversation
|
Note: This PR only fully works once diffusers have fixed this issue huggingface/diffusers#10540 by merging this PR huggingface/diffusers#10541. |
Looks like it was merged? huggingface/diffusers#10545 |
yes, but other open issues above before this PR can be merged |
|
Finished |
| transformer = FluxTransformer2DModel.from_single_file( | ||
| transformer_model_name, | ||
| #avoid loading the transformer in float32: | ||
| torch_dtype = torch.bfloat16 if weight_dtypes.prior.torch_dtype() is None else weight_dtypes.prior.torch_dtype() |
There was a problem hiding this comment.
This could now quantize the weights twice if the source is not in bf16 format. For example, if weight_dtypes.prior is set to fp8, and the source is in fp32, then this will first convert to bf16 (losing precision). Then the actual quantization to fp8 is done.
There was a problem hiding this comment.
this is model-specific code and Flux was published in bf16, so I don't think we have to worry about precision loss from fp32.
the alternative is mentioned here, #639 (comment), to let the user choose a load dtype.
The current behaviour is to load a fp8 source to fp32 and load a bf16 source to bf16. Loading fp8 sources needs 48 GB vram just for that.
new behaviour is to load both to bf16.
There was a problem hiding this comment.
But that argument is only true for the base flux model released by BFL. Fine tunes can be distributed in any other precision. There was some discussion about a possible solution in the diffusers library, but it didn't really reach any conclusion huggingface/diffusers#10679
A load d type could work, but I can almost guarantee you that there will be a lot of questions about that setting
There was a problem hiding this comment.
I'd leave it as proposed in the PR.
It's not ideal, but an improvement for all current use cases, which is fp8 and bf16 transformer files. Precision loss can happen only if you had a fp32 transformer file and wanted to load it in fp8, and what's the point of that? If you specifically created an fp32 transformer for some reason on your own (because none exist AFAIK), you would load it in fp32 and not in fp8
People have started to use and share Flux .safetensors files that only contain the Flux transformer, not the entire model pipeline. The entire pipeline is almost twice the size of the transformer only.
Examples:
Using this PR, you can load the model pipeline from
black-forest-labs/FLUX.1-dev, but override the transformer from a safetensors file.