Update raw.py

raw.py

@@ -3,9 +3,22 @@ import spaces
 import os
 from diffusers.utils import load_image
 from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
+from transformers import T5EncoderModel
+from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
+
 import gradio as gr
 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
+)
+
 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to("cuda")
 
 # Load pipeline
@@ -18,6 +31,7 @@ pipe = FluxControlNetPipeline.from_pretrained(
     controlnet=controlnet,
     torch_dtype=torch.bfloat16,
     vae=good_vae,
+    text_encoder_2=text_encoder_2_8bit,
     token=huggingface_token
 )
 adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
@@ -29,7 +43,7 @@ pipe.load_lora_weights(adapter_id2, adapter_name="real")
 pipe.set_adapters(["turbo", "real"], adapter_weights=[0.9, 0.6])
 pipe.fuse_lora(adapter_names=["turbo", "real"], lora_scale=1.0)
 pipe.unload_lora_weights()
-
+pipe.enable_xformers_memory_efficient_attention()
 # save to the Hub
 # pipe.push_to_hub("fused-t-r")
 
@@ -53,13 +67,14 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
     return image
 
 # Create Gradio interface with rows and columns
-with gr.Blocks(title="FLUX ControlNet Image Generation") as iface:
+with gr.Blocks(title="FLUX ControlNet Image Generation", fill_height=True) as iface:
     gr.Markdown("# FLUX ControlNet Image Generation")
     gr.Markdown("Generate images using the FluxControlNetPipeline. Upload a control image and enter a prompt to create an image.")
     
     with gr.Row():
         control_image = gr.Image(type="pil", label="Control Image")
         generated_image = gr.Image(type="pil", label="Generated Image", format="png")
+    with gr.Row():
         with gr.Column(scale=1):
             scale = gr.Slider(1, 3, value=1, label="Scale")
             steps = gr.Slider(6, 30, value=8, label="Steps")