Update app_v5.py

app_v5.py

@@ -10,8 +10,8 @@ import io
 import moondream as md
 from transformers import T5EncoderModel
 from diffusers import FluxControlNetPipeline, FluxPipeline, AutoModel
-from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
-from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
+from diffusers.quantizers import PipelineQuantizationConfig
+
 from diffusers.utils import load_image
 from PIL import Image
 from threading import Thread
@@ -44,25 +44,30 @@ try:
 except Exception as e:
     print(f"Error setting memory usage: {e}")
 
-quant_config_5_t5 = TransformersBitsAndBytesConfig(load_in_8bit=True,)
-text_encoder_2_8b = T5EncoderModel.from_pretrained(
-    "LPX55/FLUX.1-merged_lightning_v2",
-    subfolder="text_encoder_2",
-    quantization_config=quant_config_5_t5,
-    torch_dtype=torch.float16,
+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"],
 )
 
-quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)
-transformer_8bit = FluxPipeline.from_pretrained(
-    "LPX55/FLUX.1-merged_lightning_v2",
-    subfolder="transformer",
-    quantization_config=quant_config,
-    torch_dtype=torch.float16,
-)
+# quant_config_5_t5 = TransformersBitsAndBytesConfig(load_in_8bit=True,)
+# text_encoder_2_8b = T5EncoderModel.from_pretrained(
+#     "LPX55/FLUX.1-merged_lightning_v2",
+#     subfolder="text_encoder_2",
+#     quantization_config=quant_config_5_t5,
+#     torch_dtype=torch.float16,
+# )
+
+# quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)
+# transformer_8bit = FluxPipeline.from_pretrained(
+#     "LPX55/FLUX.1-merged_lightning_v2",
+#     subfolder="transformer",
+#     quantization_config=quant_config,
+#     torch_dtype=torch.float16,
+# )
 pipe = FluxControlNetPipeline.from_pretrained(
     "LPX55/FLUX.1M-8step_upscaler-cnet",
-    transformer=transformer_8bit,
-    text_encoder_2=text_encoder_2_8b,
+    quantization_config=pipeline_quant_config,
     torch_dtype=torch.float16,
     device_map="auto",
 )