Update app.py

app.py

@@ -29,41 +29,47 @@ 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)
 
-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)
 
+group_offloading = None
+
 # --- Inference Function ---
 @spaces.GPU
 def generate_image(prompt: str, seed: int = 42, width: int = DEFAULT_WIDTH, height: int = DEFAULT_HEIGHT, randomize_seed: bool = False, num_inference_steps: int = DEFAULT_INFERENCE_STEPS, is_enhance: bool = False):
     """Generates an image using the FLUX pipeline with error handling."""
+
+    global group_offloading
+    if not group_offloading:
+        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,
+        )
+
+        group_offloading = True
     
     if pipe is None:
         raise gr.Error("Diffusion pipeline failed to load. Cannot generate images.")