Update optimized.py

optimized.py

@@ -4,28 +4,46 @@ import os
 from diffusers.utils import load_image
 from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
 import gradio as gr
+from accelerate import init_empty_weights
+
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
-good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, use_safetensors=True, token=huggingface_token).to("cuda")
+try:
+    import xformers
+    pipe.enable_xformers_memory_efficient_attention()
+except ImportError:
+    print("XFormers missing! Using PyTorch attention instead")
+    # Fallback to PyTorch 2.0+ memory efficient attention
+    pipe.enable_sdp_attention()
+    torch.backends.cuda.enable_flash_sdp(True)
+
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", 
+                                         torch_dtype=torch.bfloat16, 
+                                         # variant="4bit",    
+                                         device_map="balanced",
+                                         use_safetensors=True, 
+                                         token=huggingface_token).to("cuda")
 
 # Load pipeline
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler",
     torch_dtype=torch.bfloat16
 )
-pipe = FluxControlNetPipeline.from_pretrained(
-    "LPX55/FLUX.1-merged_uncensored",
-    controlnet=controlnet,
-    torch_dtype=torch.bfloat16,
-    device_map="balanced",
-    vae=good_vae,
-    use_safetensors=True, 
-    token=huggingface_token
-)
+with init_empty_weights():
+    pipe = FluxControlNetPipeline.from_pretrained(
+        "LPX55/FLUX.1-merged_uncensored",
+        controlnet=controlnet,
+        torch_dtype=torch.bfloat16,
+        device_map="balanced",
+        vae=good_vae,
+        use_safetensors=True, 
+        token=huggingface_token
+    )
+pipe.enable_model_cpu_offload(device="cuda")
 # Add to your pipeline initialization:
-pipe.enable_xformers_memory_efficient_attention()
+# pipe.enable_xformers_memory_efficient_attention()
 # pipe.enable_vae_slicing()  # Batch processing of VAE
-pipe.enable_model_cpu_offload()  # Use with accelerate
+# pipe.enable_model_cpu_offload()  # Use with accelerate
 
 # Convert all models to memory-efficient format
 pipe.to(memory_format=torch.channels_last)