setup.py

import os
import sys
import re
import json
import argparse
from pathlib import Path
import subprocess
import shlex
from IPython import get_ipython

# --- Configuration ---
# You can change this to any name you want for your main folder.
WEBUI_FOLDER_NAME = 'stable-diffusion-webui-reForge'

# --- Constants for Environment Detection ---
ENV = {
    'Colab': ('/content', '/content', 'COLAB_JUPYTER_TOKEN'),
    'Kaggle': ('/kaggle', '/kaggle/working', 'KAGGLE_DATA_PROXY_TOKEN')
}

# --- Color Codes for Readable Output ---
RESET = '\033[0m'
RED = '\033[31m'
GREEN = '\033[32m'
YELLOW = '\033[33m'
ARROW = f'{YELLOW}▶{RESET}'
ERROR = f'{RED}[ERROR]{RESET}'
SUCCESS = f'{GREEN}[SUCCESS]{RESET}'
INFO = f'{YELLOW}[INFO]{RESET}'

# --- 1. Environment Detection ---
def get_environment():
    """Detects if the script is running on Colab or Kaggle and returns paths."""
    for name, (base, home, var) in ENV.items():
        if var in os.environ:
            print(f"{INFO} Detected environment: {name}")
            return name, Path(base), Path(home)
    print(f"{ERROR} This script must be run on Google Colab or Kaggle.")
    sys.exit(1)

# --- 2. Argument Parsing for API Keys ---
def get_args():
    """Parses command-line arguments for API keys."""
    parser = argparse.ArgumentParser(description='Setup API keys and directory shortcuts for model downloads.')
    parser.add_argument('--civitai_key', required=True, help='Your CivitAI API key.')
    parser.add_argument('--hf_read_token', default=None, help='Your Huggingface READ token (optional).')

    args = parser.parse_args()

    civitai_key = args.civitai_key.strip()
    hf_read_token = args.hf_read_token.strip() if args.hf_read_token else ''

    if not civitai_key:
        print(f"{ERROR} Civitai API key is missing.")
        sys.exit(1)
    if re.search(r'\s+', civitai_key):
        print(f"{ERROR} Civitai API key cannot contain spaces.")
        sys.exit(1)
    if len(civitai_key) < 32:
        print(f"{ERROR} Civitai API key seems too short (must be at least 32 characters).")
        sys.exit(1)

    print(f"{SUCCESS} Civitai API key validated.")
    if hf_read_token:
        print(f"{SUCCESS} Huggingface token provided.")

    return civitai_key, hf_read_token

# --- 3. API Key Injection ---
def download_helper_script(script_path):
    """Downloads the helper script (nenen88.py) if it doesn't exist."""
    if not script_path.exists():
        print(f"{ARROW} Downloading helper script to {script_path}...")
        script_url = 'https://github.com/gutris1/segsmaker/raw/main/script/nenen88.py'
        try:
            subprocess.run(['wget', '-qO', str(script_path), script_url], check=True)
            print(f"{SUCCESS} Helper script downloaded.")
        except subprocess.CalledProcessError:
            print(f"{ERROR} Failed to download helper script. API key injection will not work.")
            sys.exit(1)

def inject_keys(civitai_key, hf_read_token, script_path):
    """Injects API keys into the downloaded helper script."""
    print(f"{ARROW} Injecting API keys into helper script...")
    try:
        content = script_path.read_text()
        content = content.replace("TOKET = ''", f"TOKET = '{civitai_key}'")
        content = content.replace("TOBRUT = ''", f"TOBRUT = '{hf_read_token}'")
        script_path.write_text(content)
        print(f"{SUCCESS} API keys injected successfully.")
    except FileNotFoundError:
        print(f"{ERROR} Helper script not found at {script_path}. Cannot inject keys.")
        sys.exit(1)

# --- 4. Directory Structure Creation (No Shortcuts) ---
def setup_directory_structure(webui_path, model_path):
    """
    Creates a standard ReForge/A1111 directory structure.
    (Shortcuts are not created, only the directories themselves).
    """
    print(f"{ARROW} Setting up ReForge-style directory structure in {webui_path}...")

    model_dirs = {
        'Stable-diffusion',
        'VAE',
        'Lora',
        'LyCORIS',
        'embeddings',
        'ControlNet',
        'adetailer'
    }

    webui_path.mkdir(parents=True, exist_ok=True)
    model_path.mkdir(parents=True, exist_ok=True)

    for subdir in model_dirs:
        (model_path / subdir).mkdir(parents=True, exist_ok=True)
        print(f"  - Created directory: {model_path / subdir}")

    print(f"{SUCCESS} Directory structure is ready in {webui_path}")

# --- Main Execution ---
def main():
    """Main function to orchestrate the setup process."""
    env_name, base_path, home_path = get_environment()
    
    webui_path = base_path / WEBUI_FOLDER_NAME
    model_path = webui_path / 'models'
    
    startup_dir = Path('/root/.ipython/profile_default/startup')
    startup_dir.mkdir(parents=True, exist_ok=True)
    helper_script_path = startup_dir / 'nenen88.py'

    civitai_key, hf_read_token = get_args()

    download_helper_script(helper_script_path)
    inject_keys(civitai_key, hf_read_token, helper_script_path)

    # --- THE FIX: Execute the helper script to load its magic commands ---
    print(f"{ARROW} Loading downloader functions...")
    get_ipython().run_line_magic('run', str(helper_script_path))
    print(f"{SUCCESS} Downloader is now ready!")

    setup_directory_structure(webui_path, model_path)

    print("\n" + "="*50)
    print(f"{SUCCESS} Setup complete!")
    print(f"{INFO} Your main folder is: {webui_path}")
    print(f"{INFO} You can now use the %download magic command directly in your cells.")
    print(f"{INFO} Models will be saved in the 'models' subdirectory of your main folder.")
    print("="*50)


if __name__ == '__main__':
    main()