Update src/flux/pipeline_tools.py

src/flux/pipeline_tools.py

@@ -513,9 +513,7 @@ class CustomFluxPipeline:
         model_path = os.getenv("FLUX_MODEL_PATH", "black-forest-labs/FLUX.1-dev")
         print("[CustomFluxPipeline] Loading FLUX Pipeline")
         self.pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch_dtype).to(device)
-
-        # self.pipe.vae = self.pipe.vae.to(device, dtype=torch_dtype)
-        # self.pipe.transformer = self.pipe.transformer.to(device, dtype=torch_dtype)
+        self.pipe.enable_sequential_cpu_offload()
 
         self.config = config
         self.device = device