Update app.py

app.py

@@ -1,62 +1,128 @@
 import os
-import subprocess
 import shutil
+import subprocess
 import threading
 import time
 
 FOOOCUS_REPO = "https://github.com/lllyasviel/Fooocus.git"
 FOOOCUS_DIR = "Fooocus"
 
-# EPHEMERAL OUTPUTS (not saved)
+# Your direct download URL to the .safetensors file
+MODEL_URL = "PASTE_YOUR_DIRECT_DOWNLOAD_URL_HERE"
+MODEL_FILENAME = "realism_by_stable_yogi_pony.safetensors"
+
+# Save outputs to /tmp so nothing persists
 TMP_OUTPUTS = "/tmp/fooocus_outputs"
+CLEAN_EVERY_SECONDS = 60
 
-# CIVITAI MODEL (Realism by Stable Yogi – Pony)
-# You must be logged in to Civitai to get a direct file URL if required
-MODEL_URL = "PASTE_DIRECT_DOWNLOAD_URL_HERE"
-MODEL_NAME = "realism_by_stable_yogi_pony.safetensors"
 
-def run(cmd, cwd=None):
-    print(f"[RUN] {cmd}")
-    subprocess.run(cmd, shell=True, cwd=cwd, check=True)
+def sh(cmd, cwd=None, check=True):
+    print(f"\n[CMD] {cmd}\n", flush=True)
+    return subprocess.run(cmd, shell=True, cwd=cwd, check=check)
 
-def setup():
-    # Clone Fooocus
-    if not os.path.exists(FOOOCUS_DIR):
-        run(f"git clone {FOOOCUS_REPO}")
 
-    # Prepare ephemeral outputs
+def ensure_dir(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def reset_tmp_outputs():
     shutil.rmtree(TMP_OUTPUTS, ignore_errors=True)
-    os.makedirs(TMP_OUTPUTS, exist_ok=True)
+    ensure_dir(TMP_OUTPUTS)
+
 
+def link_outputs_to_tmp():
+    # Fooocus writes to Fooocus/outputs — redirect to /tmp
     outputs_path = os.path.join(FOOOCUS_DIR, "outputs")
-    if os.path.exists(outputs_path):
+
+    # remove existing outputs folder/symlink
+    if os.path.islink(outputs_path):
+        os.unlink(outputs_path)
+    elif os.path.exists(outputs_path):
         shutil.rmtree(outputs_path, ignore_errors=True)
+
     os.symlink(TMP_OUTPUTS, outputs_path)
+    print(f"[OK] outputs -> {TMP_OUTPUTS}", flush=True)
 
-    # Install dependencies once
-    marker = os.path.join(FOOOCUS_DIR, ".installed")
-    if not os.path.exists(marker):
-        run("pip install -r requirements_versions.txt", cwd=FOOOCUS_DIR)
-        open(marker, "w").close()
 
-    # Create checkpoints folder
+def download_checkpoint():
     ckpt_dir = os.path.join(FOOOCUS_DIR, "models", "checkpoints")
-    os.makedirs(ckpt_dir, exist_ok=True)
+    ensure_dir(ckpt_dir)
+
+    dst = os.path.join(ckpt_dir, MODEL_FILENAME)
+
+    if os.path.exists(dst) and os.path.getsize(dst) > 50_000_000:
+        print(f"[OK] Model already exists: {dst}", flush=True)
+        return
+
+    # Download (fails loudly if URL is wrong/403)
+    print(f"[INFO] Downloading model to: {dst}", flush=True)
+    sh(f"wget -O '{dst}' '{MODEL_URL}'", check=True)
+
+    # Sanity check: make sure it’s not HTML
+    size = os.path.getsize(dst)
+    print(f"[INFO] Downloaded size: {size} bytes", flush=True)
+    if size < 50_000_000:
+        print("[WARN] File looks too small for a checkpoint. Likely not a direct .safetensors download link or got blocked (403/HTML).", flush=True)
 
-    # Download model if missing
-    model_path = os.path.join(ckpt_dir, MODEL_NAME)
-    if not os.path.exists(model_path):
-        run(f"wget -O {model_path} {MODEL_URL}")
 
-def cleaner():
+def install_fooocus_once():
+    marker = os.path.join(FOOOCUS_DIR, ".installed")
+    if os.path.exists(marker):
+        print("[OK] Fooocus deps already installed.", flush=True)
+        return
+
+    sh("pip install -r requirements_versions.txt", cwd=FOOOCUS_DIR, check=True)
+    with open(marker, "w") as f:
+        f.write("ok\n")
+    print("[OK] Fooocus deps installed.", flush=True)
+
+
+def cleaner_loop():
     while True:
-        shutil.rmtree(TMP_OUTPUTS, ignore_errors=True)
-        os.makedirs(TMP_OUTPUTS, exist_ok=True)
-        time.sleep(60)
+        try:
+            # wipe outputs every minute
+            for name in os.listdir(TMP_OUTPUTS):
+                p = os.path.join(TMP_OUTPUTS, name)
+                if os.path.isdir(p):
+                    shutil.rmtree(p, ignore_errors=True)
+                else:
+                    try:
+                        os.remove(p)
+                    except FileNotFoundError:
+                        pass
+            print("[CLEAN] Cleared /tmp outputs", flush=True)
+        except Exception as e:
+            print(f"[CLEAN-WARN] {e}", flush=True)
+        time.sleep(CLEAN_EVERY_SECONDS)
+
+
+def main():
+    print("[BOOT] Starting Space…", flush=True)
+
+    # Clone Fooocus if missing
+    if not os.path.exists(FOOOCUS_DIR):
+        sh(f"git clone {FOOOCUS_REPO}", check=True)
+
+    # outputs -> /tmp (non-persistent)
+    reset_tmp_outputs()
+    link_outputs_to_tmp()
+
+    # install once
+    install_fooocus_once()
+
+    # download your model (only works if URL is direct and not blocked)
+    if MODEL_URL and "PASTE_" not in MODEL_URL:
+        download_checkpoint()
+    else:
+        print("[INFO] MODEL_URL not set. Skipping model download.", flush=True)
+
+    # start cleaner
+    threading.Thread(target=cleaner_loop, daemon=True).start()
+
+    # launch Fooocus
+    print("[BOOT] Launching Fooocus…", flush=True)
+    sh("python entry_with_update.py --listen 0.0.0.0 --port 7860", cwd=FOOOCUS_DIR, check=True)
 
-def start():
-    setup()
-    threading.Thread(target=cleaner, daemon=True).start()
-    run("python entry_with_update.py --listen 0.0.0.0 --port 7860", cwd=FOOOCUS_DIR)
 
-start()
+if __name__ == "__main__":
+    main()