Update app.py

app.py

@@ -1,47 +1,162 @@
+# app.py - HF Space launcher (Option B, HF tree aware)
 import os
 import subprocess
-import gradio as gr
+import threading
+import shlex
+import time
 import requests
+import gradio as gr
 
 # -----------------------------
-# Runtime downloads
+# Configuration (edit if needed)
 # -----------------------------
-downloads = {
-    "lora": "https://huggingface.co/latent-consistency/lcm-lora-sdv1-5/resolve/main/pytorch_lora_weights.safetensors",
-    "controlnet_tile": "https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11f1e_sd15_tile.pth",
-    "temporalnet": "https://huggingface.co/CiaraRowles/TemporalNet/resolve/main/diff_control_sd15_temporalnet_fp16.safetensors",
-    "civitai_model": "https://civitai.com/api/download/models/143906?type=Model&format=SafeTensor&size=pruned&fp=fp16"
+DOWNLOADS = {
+    "lora": {
+        "url": "https://huggingface.co/latent-consistency/lcm-lora-sdv1-5/resolve/main/pytorch_lora_weights.safetensors",
+        "dest": "stable-diffusion-webui/models/Lora/pytorch_lora_weights.safetensors"
+    },
+    "controlnet_tile": {
+        "url": "https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11f1e_sd15_tile.pth",
+        "dest": "stable-diffusion-webui/extensions/ControlNet/models/control_v11f1e_sd15_tile.pth"
+    },
+    "temporalnet": {
+        "url": "https://huggingface.co/CiaraRowles/TemporalNet/resolve/main/diff_control_sd15_temporalnet_fp16.safetensors",
+        "dest": "stable-diffusion-webui/extensions/ControlNet/models/diff_control_sd15_temporalnet_fp16.safetensors"
+    },
+    "civitai_model": {
+        "url": "https://civitai.com/api/download/models/143906?type=Model&format=SafeTensor&size=pruned&fp=fp16",
+        "dest": "stable-diffusion-webui/models/Stable-diffusion/civitai_model.safetensors"
+    }
 }
 
-# Ensure persistent folders exist
-os.makedirs("stable-diffusion-webui/models/Stable-diffusion", exist_ok=True)
-os.makedirs("stable-diffusion-webui/models/Lora", exist_ok=True)
-os.makedirs("stable-diffusion-webui/extensions/ControlNet/models", exist_ok=True)
-
-# Download helper
-def download_file(url, dest):
-    if not os.path.exists(dest):
-        print(f"Downloading {url} -> {dest}")
-        response = requests.get(url, stream=True)
-        with open(dest, "wb") as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
+WEBUI_ARGS = "--listen --xformers --enable-insecure-extension-access --medvram"
+LAUNCH_PY = "launch.py"  # relative to repo root (per HF tree)
+
+# -----------------------------
+# Utilities
+# -----------------------------
+def ensure_persistent_folders():
+    """Make sure HF tree persistent folders exist (safe no-op if present)."""
+    folders = [
+        "stable-diffusion-webui/deforum/input",
+        "stable-diffusion-webui/deforum/output_committed/frames",
+        "stable-diffusion-webui/deforum/output_committed/video",
+        "stable-diffusion-webui/models/Stable-diffusion",
+        "stable-diffusion-webui/models/Lora",
+        "stable-diffusion-webui/extensions/ControlNet/models",
+    ]
+    for f in 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."""
+    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:
+        try:
+            print(f"⬇️ Download attempt {attempt}: {url} -> {dest}")
+            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
+            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
+    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(f"{dest} already exists, skipping download.")
+        print("⚠️ Some runtime downloads failed. Check logs for details.")
 
-# Download models
-download_file(downloads["lora"], "stable-diffusion-webui/models/Lora/pytorch_lora_weights.safetensors")
-download_file(downloads["controlnet_tile"], "stable-diffusion-webui/extensions/ControlNet/models/control_v11f1e_sd15_tile.pth")
-download_file(downloads["temporalnet"], "stable-diffusion-webui/extensions/ControlNet/models/diff_control_sd15_temporalnet_fp16.safetensors")
-download_file(downloads["civitai_model"], "stable-diffusion-webui/models/Stable-diffusion/civitai_model.safetensors")
 
 # -----------------------------
-# Install dependencies
+# Start WebUI routine (clean)
 # -----------------------------
-subprocess.run(["pip", "install", "-r", "requirements.txt"])
+def start_webui():
+    """
+    1) Ensure folders exist
+    2) Fetch runtime models (only if missing)
+    3) Start Automatic1111's launch.py in the same foreground process context
+       using a subprocess call built from a proper arg list.
+    """
+    try:
+        ensure_persistent_folders()
 
+        # Download models (non-blocking safety: we do this before launch so they are present)
+        fetch_models()
+
+        # Build the command safely
+        args = ["python", LAUNCH_PY] + shlex.split(WEBUI_ARGS)
+        print("▶️ Launching WebUI with:", " ".join(args))
+
+        # Run launch.py and block this thread until it exits; logs go to HF logs.
+        # We run in a separate Python thread (caller keeps main process alive via Gradio).
+        subprocess.run(args, check=True)
+        print("ℹ️ launch.py exited (subprocess finished). If this happens early, check errors above.")
+    except subprocess.CalledProcessError as cpe:
+        print(f"❌ launch.py exited with non-zero status: {cpe.returncode}")
+    except Exception as e:
+        print("❌ Exception in start_webui():", e)
+
+
+# -----------------------------
+# Kick off WebUI in background thread
 # -----------------------------
-# Launch WebUI
+threading.Thread(target=start_webui, daemon=True).start()
+print("🧵 WebUI thread started (background).")
+
+# -----------------------------
+# Minimal Gradio UI to keep Space alive + quick controls
 # -----------------------------
-os.environ["COMMANDLINE_ARGS"] = "--listen --xformers --enable-insecure-extension-access --medvram"
-subprocess.run(["python", "launch.py"])
+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
+    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)
+
+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.")
+    status_btn = gr.Button("Check status")
+    status_out = gr.Textbox(lines=10)
+    status_btn.click(fn=show_status, inputs=None, outputs=status_out)
+
+demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))