Update src/pipeline.py

src/pipeline.py

@@ -17,6 +17,15 @@ import os
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False,garbage_collection_threshold:0.01"
 Pipeline = None
 
+# Define the quantization config
+nf4_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
 def empty_cache():
     start = time.time()