Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py - HF Space launcher (Option B, HF tree aware, auto-open WebUI)
+# app.py - HF Space launcher (Option B, HF tree aware, clean logs)
 import os
 import subprocess
 import threading
@@ -6,10 +6,10 @@ import shlex
 import time
 import requests
 import gradio as gr
-import webbrowser
+import socket
 
 # -----------------------------
-# Configuration (edit if needed)
+# Configuration
 # -----------------------------
 DOWNLOADS = {
     "lora": {
@@ -31,13 +31,13 @@ DOWNLOADS = {
 }
 
 WEBUI_ARGS = "--listen --xformers --enable-insecure-extension-access --medvram"
-LAUNCH_PY = "launch.py"  # relative to repo root (per HF tree)
+LAUNCH_PY = "launch.py"  # root of repo
 
 # -----------------------------
 # Utilities
 # -----------------------------
-def ensure_persistent_folders():
-    """Make sure HF tree persistent folders exist (safe no-op if present)."""
+def ensure_folders():
+    """Ensure HF tree persistent folders exist."""
     folders = [
         "stable-diffusion-webui/deforum/input",
         "stable-diffusion-webui/deforum/output_committed/frames",
@@ -50,127 +50,91 @@ def ensure_persistent_folders():
         os.makedirs(f, exist_ok=True)
     print("✅ Persistent folders ensured.")
 
-
-def download_file_with_retries(url, dest, max_retries=3, backoff=5):
-    """Download a file with simple retry/backoff. Skips if dest exists."""
+def download_file(url, dest, retries=3, backoff=5):
+    """Download a file if missing, with retries."""
     if os.path.exists(dest):
         print(f"✅ Already exists: {dest}")
         return True
-
     os.makedirs(os.path.dirname(dest), exist_ok=True)
-    attempt = 1
-    while attempt <= max_retries:
+    for attempt in range(1, retries + 1):
         try:
-            print(f"⬇️ Download attempt {attempt}: {url} -> {dest}")
+            print(f"⬇️ Download attempt {attempt}: {url}")
             with requests.get(url, stream=True, timeout=60) as r:
                 r.raise_for_status()
                 with open(dest + ".part", "wb") as f:
                     for chunk in r.iter_content(chunk_size=8192):
                         if chunk:
                             f.write(chunk)
-            # atomically move into place
             os.replace(dest + ".part", dest)
             print(f"✅ Downloaded: {dest}")
             return True
         except Exception as e:
-            print(f"⚠️ Download failed (attempt {attempt}) for {url}: {e}")
-            attempt += 1
+            print(f"⚠️ Attempt {attempt} failed: {e}")
             time.sleep(backoff * attempt)
-    print(f"❌ Failed to download after {max_retries} attempts: {url}")
-    # ensure partial file removed
-    try:
-        if os.path.exists(dest + ".part"):
-            os.remove(dest + ".part")
-    except Exception:
-        pass
+    print(f"❌ Failed to download: {dest}")
     return False
 
-
 def fetch_models():
-    """Download all configured models into the HF tree persistent paths."""
-    print("🔽 Starting runtime model downloads (if missing)...")
-    success_all = True
-    for name, info in DOWNLOADS.items():
-        url = info["url"]
-        dest = info["dest"]
-        ok = download_file_with_retries(url, dest)
-        success_all = success_all and ok
-    if success_all:
-        print("✅ All runtime downloads done (or already present).")
-    else:
-        print("⚠️ Some runtime downloads failed. Check logs for details.")
+    print("🔽 Fetching runtime models...")
+    for key, info in DOWNLOADS.items():
+        download_file(info["url"], info["dest"])
+    print("✅ Model downloads done.")
 
+# -----------------------------
+# Check if port is free
+# -----------------------------
+def wait_for_port(port, host="127.0.0.1", timeout=120):
+    """Wait until the WebUI port is open."""
+    start = time.time()
+    while time.time() - start < timeout:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            try:
+                s.connect((host, port))
+                return True
+            except:
+                time.sleep(1)
+    return False
 
 # -----------------------------
-# Start WebUI routine (clean + auto open)
+# Start WebUI thread
 # -----------------------------
 def start_webui():
-    """
-    Option B — background WebUI launcher aware of HF tree and persistent folders.
-    1) Ensure persistent folders exist.
-    2) Download missing runtime models.
-    3) Launch Automatic1111 WebUI (launch.py) with correct PORT binding.
-    4) Automatically open WebUI URL when ready.
-    """
-    try:
-        ensure_persistent_folders()
-        fetch_models()
+    ensure_folders()
+    fetch_models()
 
-        port = int(os.environ.get("PORT", 7860))
-        cmd = ["python", LAUNCH_PY] + shlex.split(WEBUI_ARGS) + [f"--port {port}"]
-        print("▶️ Launching WebUI with:", " ".join(cmd))
+    port = int(os.environ.get("PORT", 7860))
+    cmd = ["python", LAUNCH_PY] + shlex.split(WEBUI_ARGS) + [f"--port {port}"]
+    print("▶️ Launching WebUI with:", " ".join(cmd))
 
-        # Launch launch.py and capture stdout
-        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
-
-        for line in iter(proc.stdout.readline, ""):
-            print(line, end="")  # Show logs in HF Space logs
-
-            # Detect URL line
-            if "Running on local URL:" in line or "Running on" in line:
-                url_start = line.find("http")
-                if url_start != -1:
-                    url = line[url_start:].strip()
-                    print(f"🌐 WebUI ready: {url} — opening automatically")
-                    webbrowser.open(url)
-
-        proc.wait()
+    # Run launch.py once; logs go to HF logs
+    try:
+        subprocess.run(cmd, check=True)
         print("ℹ️ launch.py exited normally.")
-    except subprocess.CalledProcessError as cpe:
-        print(f"❌ launch.py exited with non-zero status: {cpe.returncode}")
+    except subprocess.CalledProcessError as e:
+        print(f"❌ launch.py failed: {e}")
     except Exception as e:
-        print("❌ Exception in start_webui():", e)
+        print(f"❌ Exception: {e}")
 
-
-# -----------------------------
-# Kick off WebUI in background thread
-# -----------------------------
 threading.Thread(target=start_webui, daemon=True).start()
-print("🧵 WebUI thread started (background).")
+print("🧵 WebUI background thread started.")
 
 # -----------------------------
-# Minimal Gradio UI to keep Space alive + quick controls
+# Gradio UI
 # -----------------------------
 def show_status():
-    msg_lines = ["✅ HF Space running (Option B launcher)."]
-    # quick presence checks
-    if os.path.exists(LAUNCH_PY):
-        msg_lines.append(f"launch.py: found")
-    else:
-        msg_lines.append(f"launch.py: MISSING at root!")
-    # list basic model presence
+    port = int(os.environ.get("PORT", 7860))
+    url = f"http://localhost:{port}"
+    lines = ["✅ HF Space running (Option B launcher)."]
+    lines.append(f"WebUI URL: [Click here]({url}) (opens in browser when ready)")
+    # Check models
     for key, info in DOWNLOADS.items():
         present = "yes" if os.path.exists(info["dest"]) else "no"
-        msg_lines.append(f"{key}: {present}")
-    return "\n".join(msg_lines)
-
+        lines.append(f"{key}: {present}")
+    return "\n".join(lines)
 
 with gr.Blocks() as demo:
     gr.Markdown("## Automatic1111 WebUI — HF Space launcher (Option B)")
-    gr.Markdown(
-        "This Space launches Automatic1111 in a background thread and keeps a small Gradio UI alive "
-        "so Hugging Face marks the app as initialized."
-    )
+    gr.Markdown("This Space launches Automatic1111 in a background thread and keeps a small Gradio UI alive.")
     status_btn = gr.Button("Check status")
     status_out = gr.Textbox(lines=10)
     status_btn.click(fn=show_status, inputs=None, outputs=status_out)