Update optimized.py

optimized.py

@@ -44,7 +44,7 @@ good_vae = AutoencoderKL.from_pretrained(
     use_safetensors=True,
     device_map=None,  # Disable automatic mapping
     token=huggingface_token
-)
+).to(device)
 # 2. Main Pipeline Initialization WITH VAE SCOPE
 pipe = FluxControlNetPipeline.from_pretrained(
     "LPX55/FLUX.1-merged_uncensored",
@@ -58,7 +58,7 @@ pipe = FluxControlNetPipeline.from_pretrained(
     device_map=None,
     token=huggingface_token  # Note corrected env var name
 )
-
+pipe.to(device)
 # 3. Strict Order for Optimization Steps
 # A. Apply CPU Offloading FIRST
 pipe.enable_sequential_cpu_offload()  # No arguments for new API
@@ -73,8 +73,8 @@ if getattr(pipe, "vae", None) is not None:
 
 pipe.enable_attention_slicing(1)
 # B. Enable Memory Optimizations
-pipe.enable_vae_tiling()
-pipe.enable_xformers_memory_efficient_attention()
+# pipe.enable_vae_tiling()
+# pipe.enable_xformers_memory_efficient_attention()
 
 # C. Unified Precision Handling
 # for comp in [pipe.unet, pipe.vae, pipe.controlnet]: