Update app.py

app.py

@@ -11,8 +11,10 @@ from pipeline_flux_cnet import FluxControlNetInpaintingPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 import subprocess
+
 from transformers import T5EncoderModel
 from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 
 subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
 
@@ -34,6 +36,20 @@ text_encoder_2_4bit = T5EncoderModel.from_pretrained(
     use_safetensors=True,
     token=HF_TOKEN
 )
+
+quant_config = DiffusersBitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+)
+
+transformer_4bit = FluxTransformer2DModel.from_pretrained(
+    "LPX55/FLUX.1-Inpainting-8step_uncensored",
+    subfolder="transformer",
+    quantization_config=quant_config,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    token=HF_TOKEN
+)
 # text_encoder_8bit = T5EncoderModel.from_pretrained(
 #     "black-forest-labs/FLUX.1-dev",
 #     subfolder="text_encoder_2",
@@ -55,7 +71,7 @@ pipe = FluxControlNetInpaintingPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     controlnet=controlnet,
     # text_encoder_2=text_encoder_8bit,
-    # transformer=pipe.transformer,
+    transformer=transformer_4bit,
     torch_dtype=torch.bfloat16, 
     # device_map="balanced",
     token=HF_TOKEN