updated the app.

app.py

@@ -31,12 +31,16 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print("Using device:", DEVICE)
 torch.backends.cudnn.benchmark = True
 
+PIPELINE=None
+
 # -----------------------------
 # Model / pipeline loading
 # -----------------------------
 @torch.no_grad()
 @spaces.GPU
 def load_pipeline_single_gpu() -> FluxKontextSliderPipeline:
+    global PIPELINE
+    
     pretrained = "black-forest-labs/FLUX.1-Kontext-dev"
 
     n_slider_layers = 4
@@ -78,7 +82,7 @@ def load_pipeline_single_gpu() -> FluxKontextSliderPipeline:
     # ------------------------------- --------------------- --------------------------- # 
 
     # Build full pipeline on CPU; no device_map sharding
-    pipeline = FluxKontextSliderPipeline.from_pretrained(
+    PIPELINE = FluxKontextSliderPipeline.from_pretrained(
         pretrained,
         transformer=transformer,
         slider_projector=slider_projector,
@@ -89,12 +93,12 @@ def load_pipeline_single_gpu() -> FluxKontextSliderPipeline:
 
     print("loading the pipeline lora weights from: {}".format(trained_models_path))
 
-    pipeline.load_lora_weights(trained_models_path)
+    PIPELINE.load_lora_weights(trained_models_path)
     print("loaded the pipeline with lora weights from: {}".format(trained_models_path)) 
-    return pipeline
-
 
-PIPELINE = load_pipeline_single_gpu()
+# Initializing the pipeline with gpu
+print("INIT pipeline with the gpu")
+load_pipeline_single_gpu()
 PIPELINE.to(DEVICE)
 print(f"[init] Pipeline loaded on {DEVICE}")
 
@@ -294,8 +298,8 @@ def _encode_prompt(prompt: str):
 # -----------------------------
 # Inference functions
 # -----------------------------
-@torch.no_grad()
 @spaces.GPU
+@torch.no_grad()
 def generate_image_stack_edits(text_prompt, n_edits, input_image):
     """
     Compute n_edits images on a single GPU for slider values in (0,1],