small changes

app.py

@@ -32,74 +32,98 @@ PIPELINE=None
 # Model / pipeline loading
 # -----------------------------
 @spaces.GPU
-def load_pipeline_single_gpu() -> FluxKontextSliderPipeline:
+def load_pipeline_single_gpu():
+    """Initialize PIPELINE inside the ZeroGPU worker and return a small status string."""
     global PIPELINE
-    
-    pretrained = "black-forest-labs/FLUX.1-Kontext-dev"
-
-    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    print("Using device:", DEVICE)
-    torch.backends.cudnn.benchmark = True
-
-    n_slider_layers = 4
-    slider_projector_out_dim = 6144
-    trained_models_path = "./model_weights/"
-    is_clip_input = True
-
-    # Load transformer fully on CPU; avoid meta tensors
-    transformer = FluxTransformer2DModelwithSliderConditioning.from_pretrained(
-        pretrained,
-        subfolder="transformer",
-        device_map=None,
-        low_cpu_mem_usage=False,
-        token=HF_TOKEN,
-    )
-    weight_dtype = transformer.dtype  # keep checkpoint dtype
+    if PIPELINE is not None:
+        print("[worker] PIPELINE already initialized; skipping.")
+        return "warm"
 
-    # Slider projector
-    if is_clip_input:
-        slider_projector = SliderProjector(
-            out_dim=slider_projector_out_dim, pe_dim=2, n_layers=n_slider_layers, is_clip_input=True
+    try:
+        # --- worker-local env & device ---
+        os.environ.pop("ACCELERATE_INIT_EMPTY_WEIGHTS", None)
+        HF_TOKEN = os.environ.get("HF_TOKEN")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print("[worker] cuda available:", torch.cuda.is_available())
+        if device == "cuda":
+            torch.backends.cudnn.benchmark = True
+
+        # --- config ---
+        pretrained = "black-forest-labs/FLUX.1-Kontext-dev"
+        n_slider_layers = 4
+        slider_projector_out_dim = 6144
+        trained_models_path = "./model_weights/"
+        is_clip_input = True
+
+        # --- validate files before loading ---
+        projector_path = os.path.join(trained_models_path, "slider_projector.pth")
+        if not os.path.isfile(projector_path):
+            return f"error: missing projector weights at {projector_path}"
+        if not os.path.isdir(trained_models_path):
+            return f"error: missing dir {trained_models_path}"
+
+        # --- transformer on CPU first ---
+        transformer = FluxTransformer2DModelwithSliderConditioning.from_pretrained(
+            pretrained,
+            subfolder="transformer",
+            device_map=None,
+            low_cpu_mem_usage=False,
+            token=HF_TOKEN,        # ok if None for public repos
+            # trust_remote_code=True,  # uncomment if this model requires it
         )
-    else:
-        slider_projector = SliderProjector_wo_clip(
-            out_dim=slider_projector_out_dim, pe_dim=2, n_layers=n_slider_layers
+        weight_dtype = transformer.dtype
+
+        # --- projector ---
+        if is_clip_input:
+            slider_projector = SliderProjector(
+                out_dim=slider_projector_out_dim, pe_dim=2, n_layers=n_slider_layers, is_clip_input=True
+            )
+        else:
+            slider_projector = SliderProjector_wo_clip(
+                out_dim=slider_projector_out_dim, pe_dim=2, n_layers=n_slider_layers
+            )
+
+        transformer.eval()
+        slider_projector.eval()
+
+        # --- load projector weights (CPU) ---
+        state_dict = torch.load(projector_path, map_location="cpu")
+        # small print (avoid dumping huge keys)
+        print("[worker] projector keys sample:", list(state_dict.keys())[:5])
+        slider_projector.load_state_dict(state_dict)
+
+        # --- build pipeline (CPU) ---
+        pipe = FluxKontextSliderPipeline.from_pretrained(
+            pretrained,
+            transformer=transformer,
+            slider_projector=slider_projector,
+            torch_dtype=weight_dtype,
+            device_map=None,
+            low_cpu_mem_usage=False,
         )
 
-    # putting both the models to infer 
-    transformer.eval()
-    slider_projector.eval()
-
-    # Load projector weights on CPU
-    slider_projector_path = os.path.join(trained_models_path, "slider_projector.pth")
-    state_dict = torch.load(slider_projector_path, map_location='cpu')
-    print("state_dict keys: {}".format(state_dict.keys()))
-
-    slider_projector.load_state_dict(state_dict)
-    print(f"loaded slider_projector from {slider_projector_path}")
-    # ------------------------------- --------------------- --------------------------- # 
-
-    # Build full pipeline on CPU; no device_map sharding
-    pipe = FluxKontextSliderPipeline.from_pretrained(
-        pretrained,
-        transformer=transformer,
-        slider_projector=slider_projector,
-        torch_dtype=weight_dtype,
-        device_map=None,
-        low_cpu_mem_usage=False,
-    )
+        # --- LoRA load (still in worker) ---
+        print("[worker] loading LoRA from:", trained_models_path)
+        pipe.load_lora_weights(trained_models_path)
 
-    print("loading the pipeline lora weights from: {}".format(trained_models_path))
+        # --- move to worker's device ---
+        pipe.to(device)
 
-    pipe.load_lora_weights(trained_models_path)
-    print("loaded the pipeline with lora weights from: {}".format(trained_models_path)) 
+        # keep in worker-global
+        PIPELINE = pipe
+        print("[worker] PIPELINE ready on", device)
+        return "ok"
 
-    pipe.to(DEVICE)
-    PIPELINE = pipe
-    
+    except Exception:
+        tb = traceback.format_exc()
+        print("[worker] exception during init:\n", tb)
+        # Return the text so you can see it in Space logs
+        return "error:\n" + tb    
+        
 # Initializing the pipeline with gpu
 print("INIT pipeline with the gpu")
-load_pipeline_single_gpu()
+status = load_pipeline_single_gpu()
+print("[main] worker init status:", status)
 
 # -----------------------------
 # Sample Images & Precomputed Results