Fixed Sometimes the dtype of the model is incorrect

[balala8]

1. Fixed Sometimes the dtype of the model is incorrect, This bug occurs often when loading weights from a local file.
2. Fixed the issue that the QuantizedTransformersModel model does not have a __call__ method, which caused an error when using the T5Encoder model in flux

What does this PR do?

1. Fixes Sometimes the dtype of the model is incorrect,
2. Fixed the issue that the QuantizedTransformersModel model does not have a call method, which caused an error when using the T5Encoder model in quantized flux

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you run all tests locally and make sure they pass.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[dacorvo]

same thing here.

[dacorvo]

This reverts the change I just made to fix a bug ...

[balala8]

Sorry, I will change it.

[balala8]

1. change default dtype from float16 to float32

[balala8]

Sorry, I will change it.

[dacorvo]

This is incorrect, as it will force model to torch.float32 if torch_dtype is not specified.
I think the correct fix would be:

if "torch_dtype" in kwargs:
    model = model.to(kwargs["torch_dtype")

[balala8]

I think this parameter must be provided.. I quantized the flux transformer model and the t5encoder model. After saving them locally, when reloading, the bias of a convolutional layer in the transformer model was always in fp32, while the entire t5encoder model was in fp32, which did not reflect the acceleration advantage, the code can be run without error. Forcing the model to be fp16 can avoid this problem.
this is i using code.
quanto and save,

class QuantizedFluxTransformer2DModel(QuantizedDiffusersModel):
    base_class = FluxTransformer2DModel

class QuantizedT5EncoderModelForCausalLM(QuantizedTransformersModel):
    auto_class = T5EncoderModel
    auto_class.from_config = auto_class._from_config
    
if __name__ == "__main__":
    bfl_repo = "black-forest-labs/FLUX.1-dev"
    dtype = torch.float16

    transformer = FluxTransformer2DModel.from_pretrained("black-forest-labs/FLUX.1-dev",subfolder="transformer", torch_dtype=dtype)
    q_transformer = QuantizedFluxTransformer2DModel.quantize(transformer, weights=qfloat8)
    q_transformer.save_pretrained("flux_transforemer_fp8_quanto")
    
    t5encoder = T5EncoderModel.from_pretrained("black-forest-labs/FLUX.1-dev",subfolder="text_encoder_2", torch_dtype=dtype)
    t5encoder = QuantizedT5EncoderModelForCausalLM.quantize(t5encoder, weights=qfloat8)
    t5encoder.save_pretrained("flux_T5Encoder_fp8_quanto")

here is load and inference code

import torch
from diffusers import FluxTransformer2DModel, FluxPipeline
from transformers import T5EncoderModel, CLIPTextModel
from optimum.quanto import freeze, qfloat8, quantize,QuantizedTransformersModel,QuantizedDiffusersModel
from huggingface_hub import login

bfl_repo = "black-forest-labs/FLUX.1-dev"
dtype = torch.float16

class QuantizedFluxTransformer2DModel(QuantizedDiffusersModel):
    base_class = FluxTransformer2DModel

class QuantizedT5EncoderModelForCausalLM(QuantizedTransformersModel):
    auto_class = T5EncoderModel
    auto_class.from_config = auto_class._from_config

transformer = QuantizedFluxTransformer2DModel.from_pretrained("./flux_transforemer_fp8_quanto", torch_dtype=dtype)
transformer.to(device="cuda")
print("transformer.dtype:",transformer.dtype)
text_encoder_2 = QuantizedT5EncoderModelForCausalLM.from_pretrained("./flux_T5Encoder_fp8_quanto", torch_dtype=dtype)
text_encoder_2.to(device="cuda")  
print("text_encoder_2.dtype:",text_encoder_2.dtype)

pipe = FluxPipeline.from_pretrained(bfl_repo, 
                                    transformer=None,
                                    text_encoder_2=None,
                                    torch_dtype=dtype)
pipe.transformer = transformer
pipe.text_encoder_2 = text_encoder_2
pipe = pipe.to(device="cuda")

pipe.enable_model_cpu_offload()

prompt = "cookie monster, yarn art style"
image = pipe(
    prompt,
    guidance_scale=3.5,
    output_type="pil",
    num_inference_steps=20,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flux-fp8-dev.png")

[github-actions]

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

[github-actions]

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