Support for new diffuser: flux1.schnell

[KoppAlexander]

@sayakpaul

Hello,

I am looking for support for saving and loading the flux1.schnell model from Blackforest.

Following your code from the "Bonus" here

Saving

from diffusers import PixArtTransformer2DModel
from optimum.quanto import QuantizedPixArtTransformer2DModel, qfloat8

model = PixArtTransformer2DModel.from_pretrained("PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", subfolder="transformer")
qmodel = QuantizedPixArtTransformer2DModel.quantize(model, weights=qfloat8)
qmodel.save_pretrained("pixart-sigma-fp8")

Loading

from optimum.quanto import QuantizedPixArtTransformer2DModel
import torch

transformer = QuantizedPixArtTransformer2DModel.from_pretrained("pixart-sigma-fp8") 
transformer.to(device="cuda", dtype=torch.float16)

I am looking for a similar option to save and load the two quantized models in this repo here,

see line 38,45,46

transformer = FluxTransformer2DModel.from_pretrained(bfl_repo, subfolder="transformer", torch_dtype=dtype, revision=revision)
quantize(transformer, weights=qfloat8)
freeze(transformer)

and 35,48,49.

text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype, revision=revision)
quantize(text_encoder_2, weights=qfloat8)
freeze(text_encoder_2)

If I'm correctly understanding this, there would need to be somethin like

from optimum.quanto import QuantizedFluxTransformer2DModel for the transformer
and
from optimum.quanto import T5EncoderModel for the text_encoder_2
to be able to save and load the quantized models for the transformer and encoder. Is that correct? Or is tehre another possibility which avoids importing the quantized version of a model from optimum.quanto?

Thank you!

@sayakpaul

[sayakpaul]

Yes, exactly, you will need that. But since the classes are simple as you can see here, you can probably start making progress already by having them implemented in your projects:

optimum-quanto/optimum/quanto/models/diffusers_models.py

Line 184 in 601dc19

class QuantizedPixArtTransformer2DModel(QuantizedDiffusersModel):

But of course, we would appreciate your contributions to quanto, as well :)

[KoppAlexander]

Thanks for the quick answer!

Did it like this:

class QuantizedFlux2DModel(QuantizedDiffusersModel):
    base_class = FluxTransformer2DModel

class QuantizedT5Model(QuantizedTransformersModel):
    auto_class = T5EncoderModel

Do you know if that's the correct way to use it? Or is an additional freeze step necessary between quantization and saving?

# quantizing
print("start loading transformer...")
transformer = FluxTransformer2DModel.from_pretrained(bfl_repo, subfolder="transformer", torch_dtype=dtype)#, revision=revision)
print("start quantizing...")
qtransformer = QuantizedFlux2DModel.quantize(transformer, weights=qint4, exclude=["proj_out", "x_embedder", "norm_out", "context_embedder"])
print("...end quantizing")

# NECESSARY?
freeze(qtransformer)

# saving
print("start saving...")
qtransformer.save_pretrained(f"{bfl_repo}/q_transformer")
print("...end saving")

# loading
qtransformer = QuantizedFlux2DModel.from_pretrained(f"{bfl_repo}/q_transformer")
qtransformer.to(device="cpu", dtype=dtype)

[sayakpaul]

After you quantize, it's not necessary to freeze as you can see it happens within the quantize() call:

optimum-quanto/optimum/quanto/models/diffusers_models.py

Line 114 in 601dc19

freeze(model)

Rest looks good to me.

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[github-actions]

This issue was closed because it has been stalled for 5 days with no activity.

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.