Update app.py

app.py

@@ -12,19 +12,29 @@ import random
 import time
 from huggingface_hub import hf_hub_download
 from diffusers import FluxTransformer2DModel, FluxPipeline
+from optimum.quanto.models import QuantizedDiffusersModel, QuantizedTransformersModel
 import safetensors.torch
 from safetensors.torch import load_file
 import gc
+from optimum.quanto.models import QuantizedDiffusersModel, QuantizedTransformersModel
+from tea_model import TeaDecoder
+from text_encoder import t5_config, T5EncoderModel, PretrainedTextEncoder
 
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
 
-
 torch.backends.cuda.matmul.allow_tf32 = True
 
-pipe = FluxPipeline.from_pretrained("John6666/fastflux-unchained-t5f16-fp8-flux", torch_dtype=torch.bfloat16)
+class Flux2DModel(QuantizedDiffusersModel):
+    base_class = FluxTransformer2DModel
+
+if __name__ == '__main__':
+    t5 = PretrainedTextEncoder(t5_config, T5EncoderModel(t5_config)).to(dtype=torch.float16)
+    t5.load_model('text_encoder_2.safetensors')
+
+pipe = FluxPipeline.from_pretrained("John6666/fastflux-unchained-t5f16-fp8-flux", torch_dtype=torch.bfloat16, text_encoder_2=t5)
 pipe.to(device="cuda", dtype=torch.bfloat16)
 
 # Load LoRAs from JSON file