Update app.py

app.py

@@ -27,6 +27,9 @@ from huggingface_hub import hf_hub_download
 from diffusers.quantizers import PipelineQuantizationConfig
 from diffusers import (FluxPriorReduxPipeline, FluxInpaintPipeline, FluxFillPipeline, FluxKontextPipeline, FluxPipeline)
 
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+
 # Login Hugging Face Hub
 HF_TOKEN = os.environ.get("HF_TOKEN")
 login(token=HF_TOKEN)
@@ -36,16 +39,10 @@ import diffusers
 dtype = torch.bfloat16
 device = "cuda:0"
 
-base_model = "black-forest-labs/FLUX.1-Krea-dev"
-
-# pipeline_quant_config = PipelineQuantizationConfig(
-#     quant_backend="bitsandbytes_4bit",
-#     quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
-#     components_to_quantize=["transformer", "text_encoder_2"],
-# )
-
-txt2img_pipe = FluxKontextPipeline.from_pretrained(base_model, torch_dtype=dtype)
-txt2img_pipe = txt2img_pipe.to(device)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
+txt2img_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
+torch.cuda.empty_cache()
 
 MAX_SEED = 2**32 - 1