Update optimized.py

optimized.py

@@ -8,14 +8,7 @@ from accelerate import init_empty_weights
 
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
-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, 
@@ -29,24 +22,31 @@ controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler",
     torch_dtype=torch.bfloat16
 )
-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
-    )
+#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_vae_slicing()  # Batch processing of VAE
 # pipe.enable_model_cpu_offload()  # Use with accelerate
-
+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)
 # Convert all models to memory-efficient format
-pipe.to(memory_format=torch.channels_last)
+#pipe.to(memory_format=torch.channels_last)
 pipe.to("cuda")
 
 @spaces.GPU