Update app.py

app.py

@@ -13,16 +13,26 @@ from transformers import AutoProcessor, pipeline, AutoModelForMaskGeneration
 from diffusers.models.attention_processor import Attention
 from dataclasses import dataclass
 from typing import Any, List, Dict, Optional, Union, Tuple
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, FluxTransformer2DModel, FluxPipeline
+from transformers import BitsAndBytesConfig as BitsAndBytesConfig, T5EncoderModel
 
 # Ensure that the minimal version of diffusers is installed
 check_min_version("0.30.2")
 HF_TOKEN = os.getenv("HF_TOKEN")
 os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'
-
+quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
+transformer_8bit = FluxTransformer2DModel.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    subfolder="transformer",
+    quantization_config=quant_config,
+    torch_dtype=torch.bfloat16,
+    token=HF_TOKEN
+)
 # Load necessary models and processors
 controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16)
 pipe = FluxControlNetInpaintingPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
+    "LPX55/FLUX.1-merged_uncensored",
+    transformer=transformer_8bit,
     controlnet=controlnet,
     torch_dtype=torch.bfloat16,
     use_safetensors=True,