vae slicing and tiling for flux

main.py

@@ -270,7 +270,11 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pip
             
             if args.cpu_offloading:
                 pipe.enable_sequential_cpu_offload()
-
+                pipe.vae.enable_slicing()
+                pipe.vae.enable_tiling()
+                
+                pipe.to(torch.float16) # casting here instead of in the pipeline constructor because doing so in the constructor loads all models into CPU memory at once
+            
         
         if args.enable_multi_apply:
             
@@ -299,11 +303,6 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pip
         )
         best_image.save(f"{save_dir}/best_image.png")
         #init_image.save(f"{save_dir}/init_image.png")
-
-        # Move the pipe back to CPU
-        if args.model != "flux":
-            # Call the function to find and move object to CPU
-            find_and_move_object_to_cpu()
         
         clear_gpu()