Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments to fight fragmentation; explicit low_cpu_mem_usage=False on TAEF1 + add accelerate to requirements to silence the warning

app.py

@@ -3,6 +3,9 @@ import sys
 import subprocess
 import tempfile
 
+# Help the allocator survive the large-activation spikes during PiD pixel-space ops
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
 import spaces
 
 
@@ -102,7 +105,9 @@ pipeline.to("cuda")
 
 print("[pid] loading TAEF1 (fast preview decoder)...", flush=True)
 from diffusers import AutoencoderTiny
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=DTYPE).to("cuda")
+taef1 = AutoencoderTiny.from_pretrained(
+    "madebyollin/taef1", torch_dtype=DTYPE, low_cpu_mem_usage=False
+).to("cuda")
 taef1.eval()
 
 def _load_pid(ckpt_type: str):

requirements.txt

@@ -1,3 +1,4 @@
+accelerate
 diffusers>=0.37.0
 transformers==4.57.1
 safetensors