Add fresh RunPod setup script

setup_runpod.py

@@ -0,0 +1,291 @@
+#!/usr/bin/env python3
+"""
+RunPod Fresh Setup Script for Qwen Image Edit
+==============================================
+Downloads everything from YOUR OWN HuggingFace repos:
+  - App code: sdfafdfsdf/Qwen-Image-Edit-App
+  - Models:   sdfafdfsdf/Phr00t-Qwen-Rapid-AIO (20GB)
+              sdfafdfsdf/Qwen-Image-Edit-2511   (16GB)
+              sdfafdfsdf/BFS-Best-Face-Swap      (1.3GB)
+              sdfafdfsdf/Qwen-Image-Edit-2511-Multiple-Angles-LoRA (282MB)
+              sdfafdfsdf/vllm-flash-attn3        (807MB)
+
+Usage:
+  1. Upload this script to your RunPod
+  2. Run: python3 setup_runpod.py
+  3. The app will auto-start at the end
+
+Set your HF token as an environment variable before running:
+  export HF_TOKEN="hf_your_token_here"
+
+Or pass it as a command line argument:
+  python3 setup_runpod.py --token hf_your_token_here
+"""
+
+import os
+import sys
+import subprocess
+import argparse
+import time
+
+# ═══════════════════════════════════════════════════════
+# CONFIGURATION — Change these if you rename your repos
+# ═══════════════════════════════════════════════════════
+
+HF_USERNAME = "sdfafdfsdf"
+
+# App code repo (small files: app.py, qwenimage/, camera_control_ui, etc.)
+APP_REPO = f"{HF_USERNAME}/Qwen-Image-Edit-App"
+
+# Model repos (large files: transformer weights, VAE, text encoder, etc.)
+MODEL_REPOS = [
+    f"{HF_USERNAME}/Phr00t-Qwen-Rapid-AIO",           # 20GB - AIO transformer
+    f"{HF_USERNAME}/Qwen-Image-Edit-2511",              # 16GB - Pipeline (VAE, text encoder)
+    f"{HF_USERNAME}/BFS-Best-Face-Swap",                 # 1.3GB - Face swap LoRA
+    f"{HF_USERNAME}/Qwen-Image-Edit-2511-Multiple-Angles-LoRA",  # 282MB - Angles LoRA
+    f"{HF_USERNAME}/vllm-flash-attn3",                   # 807MB - Flash attention kernel
+]
+
+# Directories
+WORKSPACE_DIR = "/workspace"
+APP_DIR = os.path.join(WORKSPACE_DIR, "Qwen-Image-Edit")
+VENV_DIR = os.path.join(WORKSPACE_DIR, "venv")
+CACHE_DIR = os.path.join(WORKSPACE_DIR, "cache")
+HF_HOME = os.path.join(CACHE_DIR, "huggingface")
+PIP_CACHE = os.path.join(CACHE_DIR, "pip")
+TMP_DIR = os.path.join(WORKSPACE_DIR, "tmp")
+
+
+# ═══════════════════════════════════════════════════════
+# HELPERS
+# ═══════════════════════════════════════════════════════
+
+def log(msg, level="INFO"):
+    """Print a formatted log message."""
+    icons = {"INFO": "ℹ️ ", "OK": "✅", "WARN": "⚠️ ", "ERR": "❌", "DL": "📦", "RUN": "🚀"}
+    icon = icons.get(level, "  ")
+    print(f"{icon} {msg}")
+
+
+def run(command, cwd=None, env_extra=None, check=True):
+    """Run a shell command, stream output, return exit code."""
+    log(f"Running: {command}", "RUN")
+    env = os.environ.copy()
+    env["TMPDIR"] = TMP_DIR
+    env["PIP_CACHE_DIR"] = PIP_CACHE
+    env["HF_HOME"] = HF_HOME
+    if env_extra:
+        env.update(env_extra)
+
+    proc = subprocess.Popen(
+        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+        text=True, cwd=cwd, env=env
+    )
+    for line in proc.stdout:
+        print(line, end="")
+    proc.wait()
+
+    if proc.returncode != 0 and check:
+        log(f"Command failed (exit {proc.returncode}): {command}", "ERR")
+    return proc.returncode
+
+
+def ensure_dirs():
+    """Create all necessary directories."""
+    for d in [APP_DIR, VENV_DIR, CACHE_DIR, HF_HOME, PIP_CACHE, TMP_DIR]:
+        os.makedirs(d, exist_ok=True)
+
+
+# ═══════════════════════════════════════════════════════
+# STEP 1: Download app code
+# ═══════════════════════════════════════════════════════
+
+def download_app_code(token):
+    """Download app code from your HF repo."""
+    log(f"Downloading app code from {APP_REPO}...", "DL")
+
+    pip = os.path.join(VENV_DIR, "bin", "pip")
+    python = os.path.join(VENV_DIR, "bin", "python")
+
+    # Use huggingface_hub to download
+    run(f"""{python} -c "
+from huggingface_hub import snapshot_download
+snapshot_download(
+    repo_id='{APP_REPO}',
+    repo_type='model',
+    local_dir='{APP_DIR}',
+    token='{token}'
+)
+print('App code downloaded successfully!')
+" """)
+
+    if os.path.exists(os.path.join(APP_DIR, "app.py")):
+        log("App code downloaded successfully!", "OK")
+    else:
+        log("App code download may have failed — app.py not found!", "ERR")
+        sys.exit(1)
+
+
+# ═════════════════════════════════���═════════════════════
+# STEP 2: Download heavy models
+# ═══════════════════════════════════════════════════════
+
+def download_models(token):
+    """Download all model repos (cached in HF_HOME)."""
+    python = os.path.join(VENV_DIR, "bin", "python")
+
+    for repo in MODEL_REPOS:
+        log(f"Downloading model: {repo}...", "DL")
+        run(f"""{python} -c "
+from huggingface_hub import snapshot_download
+snapshot_download(
+    repo_id='{repo}',
+    repo_type='model',
+    token='{token}'
+)
+print('Downloaded: {repo}')
+" """)
+
+    log("All models downloaded!", "OK")
+
+
+# ═══════════════════════════════════════════════════════
+# STEP 3: Set up Python venv
+# ═══════════════════════════════════════════════════════
+
+def setup_venv():
+    """Create virtual environment if it doesn't exist."""
+    if not os.path.exists(os.path.join(VENV_DIR, "bin", "python")):
+        log("Creating virtual environment...", "INFO")
+        run(f"python3 -m venv {VENV_DIR}")
+    else:
+        log("Virtual environment already exists.", "OK")
+
+    # Upgrade pip
+    pip = os.path.join(VENV_DIR, "bin", "pip")
+    run(f"{pip} install --upgrade pip")
+
+
+# ═══════════════════════════════════════════════════════
+# STEP 4: Install dependencies
+# ═══════════════════════════════════════════════════════
+
+def install_dependencies():
+    """Install Python packages from requirements.txt."""
+    pip = os.path.join(VENV_DIR, "bin", "pip")
+    req_file = os.path.join(APP_DIR, "requirements.txt")
+
+    if os.path.exists(req_file):
+        log("Installing dependencies from requirements.txt...", "DL")
+        run(f"{pip} install -r {req_file}")
+        log("Dependencies installed!", "OK")
+    else:
+        log("No requirements.txt found!", "WARN")
+
+    # Ensure huggingface_hub is available (needed for model downloads)
+    run(f"{pip} install huggingface_hub")
+
+
+# ═══════════════════════════════════════════════════════
+# STEP 5: Start the app
+# ═══════════════════════════════════════════════════════
+
+def start_app():
+    """Start the Gradio app."""
+    python = os.path.join(VENV_DIR, "bin", "python")
+    app_py = os.path.join(APP_DIR, "app.py")
+
+    if not os.path.exists(app_py):
+        log("app.py not found, cannot start!", "ERR")
+        return
+
+    log("Starting the Gradio app...", "RUN")
+
+    env_extra = {
+        "PYTHONPATH": APP_DIR,
+        "PYTHONUNBUFFERED": "1",
+    }
+
+    # Run in foreground so user can see output
+    run(f"{python} {app_py}", cwd=APP_DIR, env_extra=env_extra, check=False)
+
+
+# ═══════════════════════════════════════════════════════
+# MAIN
+# ═══════════════════════════════════════════════════════
+
+def main():
+    parser = argparse.ArgumentParser(description="Setup Qwen Image Edit on a fresh RunPod")
+    parser.add_argument("--token", type=str, default=None,
+                        help="HuggingFace token (or set HF_TOKEN env var)")
+    parser.add_argument("--skip-models", action="store_true",
+                        help="Skip model downloads (if already cached)")
+    parser.add_argument("--skip-app", action="store_true",
+                        help="Skip app code download (if already present)")
+    parser.add_argument("--skip-deps", action="store_true",
+                        help="Skip dependency installation")
+    parser.add_argument("--no-start", action="store_true",
+                        help="Don't auto-start the app after setup")
+    args = parser.parse_args()
+
+    # Get token
+    token = args.token or os.environ.get("HF_TOKEN", "")
+    if not token:
+        log("No HuggingFace token provided!", "ERR")
+        log("Set HF_TOKEN environment variable or pass --token", "INFO")
+        sys.exit(1)
+
+    print("=" * 60)
+    print("  Qwen Image Edit — Fresh RunPod Setup")
+    print("=" * 60)
+    start_time = time.time()
+
+    # Step 0: Create directories
+    log("Creating directories...")
+    ensure_dirs()
+
+    # Step 1: Set up venv (needed before downloading, since we use huggingface_hub)
+    log("Setting up Python virtual environment...")
+    setup_venv()
+
+    # Step 2: Install huggingface_hub first (needed for downloads)
+    pip = os.path.join(VENV_DIR, "bin", "pip")
+    run(f"{pip} install huggingface_hub")
+
+    # Step 3: Download app code
+    if not args.skip_app:
+        download_app_code(token)
+    else:
+        log("Skipping app code download (--skip-app)", "INFO")
+
+    # Step 4: Install full dependencies
+    if not args.skip_deps:
+        install_dependencies()
+    else:
+        log("Skipping dependency installation (--skip-deps)", "INFO")
+
+    # Step 5: Download models
+    if not args.skip_models:
+        download_models(token)
+    else:
+        log("Skipping model downloads (--skip-models)", "INFO")
+
+    elapsed = time.time() - start_time
+    minutes = int(elapsed // 60)
+    seconds = int(elapsed % 60)
+
+    print("")
+    print("=" * 60)
+    log(f"Setup completed in {minutes}m {seconds}s!", "OK")
+    print("=" * 60)
+
+    # Step 6: Start the app
+    if not args.no_start:
+        start_app()
+    else:
+        log("Skipping app start (--no-start). To start manually:", "INFO")
+        log(f"  bash /workspace/start_app.sh", "INFO")
+
+
+if __name__ == "__main__":
+    main()