Upload fooocus.py

fooocus.py

@@ -0,0 +1,75 @@
+import os
+import sys
+import shutil
+import subprocess
+import signal
+from urllib.request import Request, urlopen
+from tqdm import tqdm  # Progress bar library
+
+# Constants
+MODEL_URL = "https://huggingface.co/oieieio/juggernautXL_v8Rundiffusion/resolve/main/juggernautXL_v8Rundiffusion.safetensors"
+MODEL_PATH = "models/checkpoints/juggernautXL_v8Rundiffusion.safetensors"
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+
+# Cleanup temporary directory
+def cleanup_temp_dir():
+    temp_dir = "/tmp/fooocus"
+    print(f"[Cleanup] Attempting to delete content of temp dir {temp_dir}")
+    try:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        print("[Cleanup] Cleanup successful")
+    except Exception as e:
+        print(f"[Cleanup] Failed to delete content of temp dir: {e}")
+
+# Download the model with progress
+def download_file(url, target_path):
+    if os.path.exists(target_path):
+        print(f"[Model Download] Model already exists at {target_path}")
+        return
+
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+    print(f"[Model Download] Downloading: \"{url}\" to {target_path}")
+
+    headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"} if HUGGINGFACE_TOKEN else {}
+    req = Request(url, headers=headers)
+
+    try:
+        with urlopen(req) as response, open(target_path, "wb") as out_file:
+            total_size = int(response.getheader("Content-Length", 0))
+            with tqdm(total=total_size, unit="B", unit_scale=True, desc="Downloading") as progress_bar:
+                for chunk in iter(lambda: response.read(1024 * 8), b""):
+                    out_file.write(chunk)
+                    progress_bar.update(len(chunk))
+        print(f"[Model Download] Successfully downloaded model to {target_path}")
+    except Exception as e:
+        print(f"[Model Download] Failed to download {url}. Error: {e}")
+        sys.exit(1)
+
+# Launch Fooocus application
+def launch_fooocus():
+    print("[Fooocus] Launching application...")
+    try:
+        subprocess.run([sys.executable, "launch.py"], check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"[Fooocus] Failed to launch application: {e}")
+        sys.exit(1)
+
+# Signal handling for graceful shutdown
+def signal_handler(signum, frame):
+    print("[Fooocus] Shutting down gracefully...")
+    sys.exit(0)
+
+signal.signal(signal.SIGINT, signal_handler)
+signal.signal(signal.SIGTERM, signal_handler)
+
+# Main logic
+def main():
+    print(f"Python {sys.version}")
+    print("Fooocus version: 2.5.5")
+
+    cleanup_temp_dir()
+    download_file(MODEL_URL, MODEL_PATH)
+    launch_fooocus()
+
+if __name__ == "__main__":
+    main()