Update optimized.py

optimized.py

@@ -12,7 +12,20 @@ def self_attention_slicing(module, slice_size=3):
     """Modified from Diffusers' original for Flux compatibility"""
     def sliced_attention(*args, **kwargs):
         return module(*args, **kwargs)  # Remove dummy implementation <source_id data="pipeline_flux_controlnet.py" />
-
+        
+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,
+#     device_map=None,  # Disable automatic mapping
+#     token=huggingface_token
+# )
+controlnet = FluxControlNetModel.from_pretrained(
+    "jasperai/Flux.1-dev-Controlnet-Upscaler",
+    torch_dtype=torch.bfloat16
+)
 # Device management - critical fix
 pipe = FluxControlNetPipeline.from_pretrained(
     "LPX55/FLUX.1-merged_uncensored",