Update optimized.py

optimized.py

@@ -5,11 +5,18 @@ import diffusers
 import PIL
 from diffusers.utils import load_image
 from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
+from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
 import gradio as gr
 from accelerate import dispatch_model, infer_auto_device_map
 from PIL import Image
+from diffusers import FluxTransformer2DModel
+from transformers import T5EncoderModel
 import gc
 # Corrected and optimized FluxControlNet implementation
+huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+device = "cuda"
+torch_dtype = torch.bfloat16
 
 def self_attention_slicing(module, slice_size=3):
     """Modified from Diffusers' original for Flux compatibility"""
@@ -35,9 +42,23 @@ def self_attention_slicing(module, slice_size=3):
         
         return output
     return sliced_attention
-device = "cuda"
     
-huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)
+text_encoder_2_8bit = T5EncoderModel.from_pretrained(
+    "LPX55/FLUX.1-merged_uncensored",
+    subfolder="text_encoder_2",
+    quantization_config=quant_config,
+    torch_dtype=torch.bfloat16,
+    token=huggingface_token
+)
+quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)
+transformer_8bit = FluxTransformer2DModel.from_pretrained(
+    "LPX55/FLUX.1-merged_uncensored",
+    subfolder="transformer",
+    quantization_config=quant_config,
+    torch_dtype=torch.bfloat16,
+    token=huggingface_token
+)
 good_vae = AutoencoderKL.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     subfolder="vae",
@@ -54,6 +75,8 @@ pipe = FluxControlNetPipeline.from_pretrained(
         torch_dtype=torch.bfloat16
     ),
     vae=good_vae,  # Now defined in scope
+    transformer=transformer_8bit,
+    text_encoder_2=text_encoder_2_8bit,
     torch_dtype=torch.bfloat16,
     use_safetensors=True,
     device_map=None,