Update app.py

app.py

@@ -271,6 +271,88 @@ def _reset_hydra():
     except Exception:
         pass
 
+# =============================================================================
+# CHAPTER 3A: STARTUP CLEANUP (safe, configurable)
+# =============================================================================
+# Controls:
+#   CLEAN_ON_START = off | light | deep   (default: light)
+#     off   -> do nothing
+#     light -> clear ./temp, ./outputs (incl. ./outputs/backgrounds)
+#     deep  -> light + /tmp/gradio + HF caches (NOT model hub unless CLEAR_MODELS=1)
+#   CLEAR_MODELS = 0|1  (deep-only) also clears ~/.cache/huggingface/hub and ./checkpoints
+#   CLEAR_PLATFORM_TMP = 0|1 (additionally clear /tmp/gradio even in light)
+CLEAN_ON_START      = os.environ.get("CLEAN_ON_START", "light").lower().strip()
+CLEAR_MODELS        = os.environ.get("CLEAR_MODELS", "0") == "1"
+CLEAR_PLATFORM_TMP  = os.environ.get("CLEAR_PLATFORM_TMP", "0") == "1"
+
+def _safe_rmtree(path: Path) -> bool:
+    """Delete a directory only if it is our project dir or an allowed /tmp path."""
+    try:
+        path = path.resolve()
+        if not path.exists():
+            return False
+        # allow: anything under project BASE_DIR
+        if path.is_dir() and path.is_relative_to(BASE_DIR):
+            shutil.rmtree(path, ignore_errors=True)
+            return True
+        # allow: /tmp/gradio* when platform uses temp uploads
+        if str(path).startswith("/tmp/gradio"):
+            shutil.rmtree(path, ignore_errors=True)
+            return True
+    except Exception:
+        pass
+    return False
+
+def startup_cleanup():
+    """Clear caches/files so each run starts fresh."""
+    if CLEAN_ON_START not in ("light", "deep"):
+        logger.info(f"Startup cleanup skipped (CLEAN_ON_START={CLEAN_ON_START})")
+        return
+
+    to_clear: List[Path] = []
+
+    # Always clear our temp/output trees
+    to_clear += [TEMP_DIR, OUT_DIR, BACKGROUND_DIR]
+
+    # /tmp/gradio holds uploads & intermediates on many hosts
+    if CLEAN_ON_START == "deep" or CLEAR_PLATFORM_TMP:
+        gradio_tmp = Path(os.environ.get("GRADIO_TEMP_DIR", "/tmp/gradio"))
+        to_clear += [gradio_tmp]
+
+    if CLEAN_ON_START == "deep":
+        # Clear non-model HF caches
+        hf_base = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
+        to_clear += [hf_base / "diffusers", hf_base / "datasets"]
+        # Optionally clear model hub + checkpoints (will redownload!)
+        if CLEAR_MODELS:
+            to_clear += [hf_base / "hub", CHECKPOINTS]
+
+    # Perform deletions
+    cleared = []
+    for d in to_clear:
+        if isinstance(d, Path) and d.exists():
+            if _safe_rmtree(d):
+                cleared.append(str(d))
+
+    # Re-create our working directories
+    for p in (CHECKPOINTS, TEMP_DIR, OUT_DIR, BACKGROUND_DIR):
+        p.mkdir(parents=True, exist_ok=True)
+
+    # Free CUDA cache too
+    if TORCH_AVAILABLE and CUDA_AVAILABLE:
+        try:
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
+
+    logger.info(
+        f"Startup cleanup done (mode={CLEAN_ON_START}, clear_models={int(CLEAR_MODELS)}). "
+        f"Cleared: {', '.join(cleared) if cleared else 'nothing'}"
+    )
+
+
+        
+
 # =============================================================================
 # CHAPTER 4: MEMORY MANAGER
 # =============================================================================