changes for initializing the pipeline outside the inference and calling it with decorator.

app.py

@@ -27,7 +27,6 @@ if HF_TOKEN:
 # -----------------------------
 # Avoid meta-tensor init from environment leftovers
 os.environ.pop("ACCELERATE_INIT_EMPTY_WEIGHTS", None)
-PIPELINE=None
 
 # -----------------------------
 # Model / pipeline loading
@@ -35,101 +34,181 @@ PIPELINE=None
 
 def _log(msg): print(msg, flush=True)
 
-def load_pipeline_single_gpu():
-    global PIPELINE
-    if PIPELINE is not None:
-        _log("[worker] PIPELINE already initialized; skipping.")
-        return "warm"
+# def load_pipeline_single_gpu():
+#     global PIPELINE
+#     if PIPELINE is not None:
+#         _log("[worker] PIPELINE already initialized; skipping.")
+#         return "warm"
+
+#     try:
+#         os.environ.pop("ACCELERATE_INIT_EMPTY_WEIGHTS", None)
+#         token = os.environ.get("HF_TOKEN")
+#         cuda_ok = torch.cuda.is_available()
+#         _log(f"[worker] cuda available: {cuda_ok}")
+#         if cuda_ok:
+#             torch.backends.cudnn.benchmark = True
+
+#         # ---------- config ----------
+#         pretrained = "black-forest-labs/FLUX.1-Kontext-dev"
+#         trained_models_path = "./model_weights/"
+#         projector_path = os.path.join(trained_models_path, "slider_projector.pth")
+#         offload_dir = "/tmp/offload"; os.makedirs(offload_dir, exist_ok=True)
+
+#         if not os.path.isdir(trained_models_path):
+#             return f"error: missing dir {trained_models_path}"
+#         if not os.path.isfile(projector_path):
+#             return f"error: missing projector weights at {projector_path}"
+
+#         # dtype selection to cut memory
+#         if cuda_ok and torch.cuda.get_device_capability(0)[0] >= 8:
+#             dtype = torch.bfloat16
+#         elif cuda_ok:
+#             dtype = torch.float16
+#         else:
+#             dtype = torch.float32
+
+#         max_memory = {"cuda": "80GiB", "cpu": "60GiB"}  # tune if needed
+
+#         _log("[worker] loading transformer (sharded/offloaded)…")
+#         transformer = FluxTransformer2DModelwithSliderConditioning.from_pretrained(
+#             pretrained,
+#             subfolder="transformer",
+#             token=token,
+#             trust_remote_code=True,
+#             torch_dtype=dtype,
+#             low_cpu_mem_usage=True,
+#             # device_map="balanced_low_0",
+#             offload_folder=offload_dir,
+#             offload_state_dict=True,
+#             # max_memory=max_memory,
+#         )
+#         weight_dtype = transformer.dtype
+#         _log(f"[worker] transformer loaded, dtype={weight_dtype}")
+
+#         _log("[worker] building slider projector…")
+#         slider_projector = SliderProjector(out_dim=6144, pe_dim=2, n_layers=4, is_clip_input=True)
+#         slider_projector.eval()
+#         _log("[worker] loading projector weights…")
+#         state_dict = torch.load(projector_path, map_location="cpu", weights_only=True)
+#         slider_projector.load_state_dict(state_dict, strict=True)
+
+#         _log("[worker] assembling pipeline (sharded/offloaded)…")
+#         pipe = FluxKontextSliderPipeline.from_pretrained(
+#             pretrained,
+#             token=token,
+#             trust_remote_code=True,
+#             transformer=transformer,
+#             slider_projector=slider_projector,
+#             torch_dtype=weight_dtype,
+#             low_cpu_mem_usage=True,
+#             # device_map="balanced_low_0",
+#             offload_folder=offload_dir,
+#             offload_state_dict=True,
+#             # max_memory=max_memory,
+#         )
+#         _log("[worker] pipeline assembled.")
+
+#         _log(f"[worker] loading LoRA from: {trained_models_path}")
+#         pipe.load_lora_weights(trained_models_path)
+#         _log("[worker] LoRA loaded.")
+
+#         # DO NOT pipe.to("cuda") here; keep auto device_map to avoid OOM
+#         PIPELINE = pipe
+#         if cuda_ok:
+#             free, total = torch.cuda.mem_get_info()
+#             _log(f"[worker] VRAM free/total: {free/1e9:.2f}/{total/1e9:.2f} GB")
+#         _log("[worker] PIPELINE ready.")
+#         return "ok"
+
+#     except Exception:
+#         _log("[worker] init exception:\n" + traceback.format_exc())
+#         return "error"
 
-    try:
-        os.environ.pop("ACCELERATE_INIT_EMPTY_WEIGHTS", None)
-        token = os.environ.get("HF_TOKEN")
-        cuda_ok = torch.cuda.is_available()
-        _log(f"[worker] cuda available: {cuda_ok}")
-        if cuda_ok:
-            torch.backends.cudnn.benchmark = True
-
-        # ---------- config ----------
-        pretrained = "black-forest-labs/FLUX.1-Kontext-dev"
-        trained_models_path = "./model_weights/"
-        projector_path = os.path.join(trained_models_path, "slider_projector.pth")
-        offload_dir = "/tmp/offload"; os.makedirs(offload_dir, exist_ok=True)
-
-        if not os.path.isdir(trained_models_path):
-            return f"error: missing dir {trained_models_path}"
-        if not os.path.isfile(projector_path):
-            return f"error: missing projector weights at {projector_path}"
-
-        # dtype selection to cut memory
-        if cuda_ok and torch.cuda.get_device_capability(0)[0] >= 8:
-            dtype = torch.bfloat16
-        elif cuda_ok:
-            dtype = torch.float16
-        else:
-            dtype = torch.float32
-
-        max_memory = {"cuda": "80GiB", "cpu": "60GiB"}  # tune if needed
-
-        _log("[worker] loading transformer (sharded/offloaded)…")
-        transformer = FluxTransformer2DModelwithSliderConditioning.from_pretrained(
-            pretrained,
-            subfolder="transformer",
-            token=token,
-            trust_remote_code=True,
-            torch_dtype=dtype,
-            low_cpu_mem_usage=True,
-            # device_map="balanced_low_0",
-            offload_folder=offload_dir,
-            offload_state_dict=True,
-            # max_memory=max_memory,
-        )
-        weight_dtype = transformer.dtype
-        _log(f"[worker] transformer loaded, dtype={weight_dtype}")
-
-        _log("[worker] building slider projector…")
-        slider_projector = SliderProjector(out_dim=6144, pe_dim=2, n_layers=4, is_clip_input=True)
-        slider_projector.eval()
-        _log("[worker] loading projector weights…")
-        state_dict = torch.load(projector_path, map_location="cpu", weights_only=True)
-        slider_projector.load_state_dict(state_dict, strict=True)
-
-        _log("[worker] assembling pipeline (sharded/offloaded)…")
-        pipe = FluxKontextSliderPipeline.from_pretrained(
-            pretrained,
-            token=token,
-            trust_remote_code=True,
-            transformer=transformer,
-            slider_projector=slider_projector,
-            torch_dtype=weight_dtype,
-            low_cpu_mem_usage=True,
-            # device_map="balanced_low_0",
-            offload_folder=offload_dir,
-            offload_state_dict=True,
-            # max_memory=max_memory,
-        )
-        _log("[worker] pipeline assembled.")
-
-        _log(f"[worker] loading LoRA from: {trained_models_path}")
-        pipe.load_lora_weights(trained_models_path)
-        _log("[worker] LoRA loaded.")
 
-        # DO NOT pipe.to("cuda") here; keep auto device_map to avoid OOM
-        PIPELINE = pipe
-        if cuda_ok:
-            free, total = torch.cuda.mem_get_info()
-            _log(f"[worker] VRAM free/total: {free/1e9:.2f}/{total/1e9:.2f} GB")
-        _log("[worker] PIPELINE ready.")
-        return "ok"
 
-    except Exception:
-        _log("[worker] init exception:\n" + traceback.format_exc())
-        return "error"
 
+# -----------------------------
+# Loading the pipeline without any function so that it will be called directly in the inference
+# ----------------------------- 
+os.environ.pop("ACCELERATE_INIT_EMPTY_WEIGHTS", None)
+token = os.environ.get("HF_TOKEN")
+cuda_ok = torch.cuda.is_available()
+_log(f"[worker] cuda available: {cuda_ok}")
+if cuda_ok:
+    torch.backends.cudnn.benchmark = True
+
+# ---------- config ----------
+pretrained = "black-forest-labs/FLUX.1-Kontext-dev"
+trained_models_path = "./model_weights/"
+projector_path = os.path.join(trained_models_path, "slider_projector.pth")
+offload_dir = "/tmp/offload"; os.makedirs(offload_dir, exist_ok=True)
+
+# dtype selection to cut memory
+if cuda_ok and torch.cuda.get_device_capability(0)[0] >= 8:
+    dtype = torch.bfloat16
+elif cuda_ok:
+    dtype = torch.float16
+else:
+    dtype = torch.float32
+
+max_memory = {"cuda": "80GiB", "cpu": "60GiB"}  # tune if needed
+
+_log("[worker] loading transformer (sharded/offloaded)…")
+transformer = FluxTransformer2DModelwithSliderConditioning.from_pretrained(
+    pretrained,
+    subfolder="transformer",
+    token=token,
+    trust_remote_code=True,
+    torch_dtype=dtype,
+    low_cpu_mem_usage=True,
+    # device_map="balanced_low_0",
+    offload_folder=offload_dir,
+    offload_state_dict=True,
+    # max_memory=max_memory,
+)
+weight_dtype = transformer.dtype
+_log(f"[worker] transformer loaded, dtype={weight_dtype}")
+
+_log("[worker] building slider projector…")
+slider_projector = SliderProjector(out_dim=6144, pe_dim=2, n_layers=4, is_clip_input=True)
+slider_projector.eval()
+_log("[worker] loading projector weights…")
+state_dict = torch.load(projector_path, map_location="cpu", weights_only=True)
+slider_projector.load_state_dict(state_dict, strict=True)
+
+_log("[worker] assembling pipeline (sharded/offloaded)…")
+pipe = FluxKontextSliderPipeline.from_pretrained(
+    pretrained,
+    token=token,
+    trust_remote_code=True,
+    transformer=transformer,
+    slider_projector=slider_projector,
+    torch_dtype=weight_dtype,
+    low_cpu_mem_usage=True,
+    # device_map="balanced_low_0",
+    offload_folder=offload_dir,
+    offload_state_dict=True,
+    # max_memory=max_memory,
+)
+_log("[worker] pipeline assembled.")
+
+_log(f"[worker] loading LoRA from: {trained_models_path}")
+pipe.load_lora_weights(trained_models_path)
+_log("[worker] LoRA loaded.")
+
+# DO NOT pipe.to("cuda") here; keep auto device_map to avoid OOM
+PIPELINE = pipe
+if cuda_ok:
+    free, total = torch.cuda.mem_get_info()
+    _log(f"[worker] VRAM free/total: {free/1e9:.2f}/{total/1e9:.2f} GB")
+_log("[worker] PIPELINE ready.")
+
+# moving the pipeline to GPU
+PIPELINE.to('cuda')
 
 # -----------------------------
 # Sample Images & Precomputed Results
 # -----------------------------
-
 def create_sample_entry(name, image_filename, prompt, result_folder, num_results=5, result_pattern="image_{i}.png", precomputed_base="./sample_images/precomputed"):
     """
     Helper function to create a sample entry with subfolder organization.
@@ -314,7 +393,7 @@ def resize_image(img: Image.Image, target: int = 512) -> Image.Image:
 # -----------------------------
 # Inference functions
 # -----------------------------
-@spaces.GPU
+@spaces.GPU(duration=500)
 @torch.no_grad()
 def generate_image_stack_edits(text_prompt, n_edits, input_image):
     """
@@ -323,13 +402,7 @@ def generate_image_stack_edits(text_prompt, n_edits, input_image):
     """
     DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
     
-    # if pipeline is null will initialize it simply. 
-    global PIPELINE
-    if PIPELINE is None:
-        status = load_pipeline_single_gpu()
-
-    print("loaded pipeline status: {}".format(status))
-    
+    # pipelien will be loaded already in the global context and will be called here 
     if not input_image or not text_prompt or text_prompt.startswith("Please select"):
         return [], None
 
@@ -376,7 +449,7 @@ def generate_image_stack_edits(text_prompt, n_edits, input_image):
     first = results[0] if results else None
     return results, first
 
-@spaces.GPU
+@spaces.GPU(duration=80)
 def generate_single_image(text_prompt, slider_value, input_image):
     if not input_image or not text_prompt or text_prompt.startswith("Please select"):
         return None