Update app.py

app.py

@@ -6,6 +6,7 @@ import torch
 import time
 from diffusers import DiffusionPipeline, AutoencoderTiny
 from custom_pipeline import FluxWithCFGPipeline
+from diffusers.hooks import apply_group_offloading
 
 # --- Torch Optimizations ---
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -28,8 +29,34 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = FluxWithCFGPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype)
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
 
-pipe.vae.enable_slicing()
-pipe.vae.enable_tiling()
+apply_group_offloading(
+    pipe.transformer,
+    offload_type="leaf_level",
+    offload_device=torch.device("cpu"),
+    onload_device=torch.device("cuda"),
+    use_stream=True,
+)
+apply_group_offloading(
+    pipe.text_encoder, 
+    offload_device=torch.device("cpu"),
+    onload_device=torch.device("cuda"),
+    offload_type="leaf_level",
+    use_stream=True,
+)
+apply_group_offloading(
+    pipe.text_encoder_2, 
+    offload_device=torch.device("cpu"),
+    onload_device=torch.device("cuda"),
+    offload_type="leaf_level",
+    use_stream=True,
+)
+apply_group_offloading(
+    pipe.vae, 
+    offload_device=torch.device("cpu"),
+    onload_device=torch.device("cuda"),
+    offload_type="leaf_level",
+    use_stream=True,
+)
 
 pipe.to(device)