File size: 3,863 Bytes
1228919
 
41c316f
1228919
 
 
 
 
 
41c316f
 
 
 
 
1228919
0adb680
1228919
 
41c316f
 
 
1228919
 
41c316f
 
 
 
 
1228919
41c316f
 
1228919
41c316f
 
1228919
41c316f
 
 
 
 
1228919
41c316f
1228919
41c316f
1228919
 
41c316f
1228919
41c316f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1228919
 
41c316f
 
 
 
 
 
 
 
 
 
 
 
 
1228919
41c316f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1228919
 
41c316f
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import os
import shutil
import subprocess
import threading
import time

FOOOCUS_REPO = "https://github.com/lllyasviel/Fooocus.git"
FOOOCUS_DIR = "Fooocus"

# 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 = 300


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 ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def reset_tmp_outputs():
    shutil.rmtree(TMP_OUTPUTS, ignore_errors=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")

    # 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)


def download_checkpoint():
    ckpt_dir = os.path.join(FOOOCUS_DIR, "models", "checkpoints")
    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)


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:
        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)


if __name__ == "__main__":
    main()