Update src/pipeline.py

src/pipeline.py

@@ -25,7 +25,6 @@ def empty_cache():
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()
-    print(f"Flush took: {time.time() - start}")
 
 def load_pipeline() -> Pipeline:    
     empty_cache()
@@ -39,11 +38,13 @@ def load_pipeline() -> Pipeline:
     text_encoder = CLIPTextModel.from_pretrained(
         ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
     )
-    ############ Text Encoder 2 ############
     text_encoder_2 = T5EncoderModel.from_pretrained(
-        "city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16
+        "HighCWu/FLUX.1-dev-4bit",
+        subfolder="text_encoder_2",
+        torch_dtype=torch.bfloat16,
     )
 
+
     model = FluxTransformer2DModel.from_pretrained(
         "/home/sandbox/.cache/huggingface/hub/models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a", torch_dtype=dtype, use_safetensors=False
     )