Update optimized.py

optimized.py

@@ -22,48 +22,40 @@ huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 #     device_map=None,  # Disable automatic mapping
 #     token=huggingface_token
 # )
-controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler",
-    torch_dtype=torch.bfloat16
+
+good_vae = AutoencoderKL.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    subfolder="vae",
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    token=huggingface_token  # Fix typo in variable name
 )
-# Device management - critical fix
+
+# 2. Main Pipeline Initialization WITH VAE SCOPE
 pipe = FluxControlNetPipeline.from_pretrained(
     "LPX55/FLUX.1-merged_uncensored",
-    controlnet=controlnet,
-    vae=good_vae,
+    controlnet=FluxControlNetModel.from_pretrained(
+        "jasperai/Flux.1-dev-Controlnet-Upscaler",
+        torch_dtype=torch.bfloat16
+    ),
+    vae=good_vae,  # Now defined in scope
     torch_dtype=torch.bfloat16,
     use_safetensors=True,
     device_map=None,
-    token=huggingface_token
+    token=huggingface_token  # Note corrected env var name
 )
 
-# Sequence verified for Diffusers 0.20.0+
-device_map = infer_auto_device_map(
-    pipe,
-    max_memory={0:"37GiB", "cpu":"60GiB"},
-    device_types=["cuda", "cpu"]
-)
-pipe = dispatch_model(pipe, device_map=device_map, main_device="cuda")
-pipe.enable_sequential_cpu_offload()  # No arguments for new API
-
-# Precision alignment (AFTER offloading) <source_id data="pipeline_flux_controlnet.py" />
-pipe.unet.to(dtype=torch.bfloat16)
-pipe.controlnet.to(dtype=torch.bfloat16)
-pipe.vae.to(dtype=torch.bfloat16)
-
-# XFormers with Flux compatibility
-if torch.cuda.is_available():
-    try:
-        pipe.enable_xformers_memory_efficient_attention(
-            attention_op=None  # Auto-select best operator
-        )
-    except Exception as e:
-        print(f"xFormers error: {e}")
-        pipe.enable_sdp_attention(mode="math")
+# 3. Strict Order for Optimization Steps
+# A. Apply CPU Offloading FIRST
+pipe.enable_sequential_cpu_offload()
 
-# Memory format optimization 
-pipe.to(memory_format=torch.channels_last)
+# B. Enable Memory Optimizations
+pipe.enable_vae_tiling()
+pipe.enable_xformers_memory_efficient_attention()
 
+# C. Unified Precision Handling
+for comp in [pipe.unet, pipe.vae, pipe.controlnet]:
+    comp.to(dtype=torch.bfloat16)
 
 print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
 @spaces.GPU