-

app.py

@@ -160,10 +160,10 @@ transformer = FluxTransformer2DModelwithSliderConditioning.from_pretrained(
     token=token,
     trust_remote_code=True,
     torch_dtype=dtype,
-    low_cpu_mem_usage=True,
+    # low_cpu_mem_usage=True,
     # device_map="balanced_low_0",
-    offload_folder=offload_dir,
-    offload_state_dict=True,
+    # offload_folder=offload_dir,
+    # offload_state_dict=True,
     # max_memory=max_memory,
 )
 weight_dtype = transformer.dtype
@@ -184,10 +184,10 @@ PIPELINE = FluxKontextSliderPipeline.from_pretrained(
     transformer=transformer,
     slider_projector=slider_projector,
     torch_dtype=weight_dtype,
-    low_cpu_mem_usage=True,
+    # low_cpu_mem_usage=True,
     # device_map="balanced_low_0",
-    offload_folder=offload_dir,
-    offload_state_dict=True,
+    # offload_folder=offload_dir,
+    # offload_state_dict=True,
     # max_memory=max_memory,
 )
 _log("[worker] pipeline assembled.")