Update src/pipeline.py

src/pipeline.py

@@ -9,8 +9,7 @@ from diffusers import FluxTransformer2DModel, DiffusionPipeline
 import gc
 import os
 from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
-import torch._dynamo
-from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+
 
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 HOME = os.environ["HOME"]
@@ -26,14 +25,10 @@ def empty_cache():
     
 def load_pipeline() -> Pipeline:
     empty_cache()
-    text_encoder = CLIPTextModel.from_pretrained(ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16)
-    quantize_(text_encoder, int8_weight_only())
     vae = AutoencoderTiny.from_pretrained("aifeifei798/taef1", torch_dtype=torch.bfloat16)
-    quantize_(vae, int8_weight_only())
+    model = FluxTransformer2DModel.from_pretrained(f"{HOME}/.cache/huggingface/hub/models--slobers--transgender/snapshots/cb99836efa0ed55856970269c42fafdaa0e44c5d", torch_dtype=torch.bfloat16, use_safetensors=False)
     text_encoder_2 = T5EncoderModel.from_pretrained("city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16)
-    quantize_(text_encoder_2, int8_weight_only())
-    model = FluxTransformer2DModel.from_pretrained(ckpt_id, subfolder="transformer", torch_dtype=torch.bfloat16, use_safetensors=False)
-    pipeline = DiffusionPipeline.from_pretrained(ckpt_id, text_encoder=text_encoder, transformer=model, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16)
+    pipeline = DiffusionPipeline.from_pretrained(ckpt_id, vae=vae, transformer=model, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
     
     for _ in range(2):