CivitAI Downloader Script

civitai-downloader.py

import os
import sys
import re
import typer
import requests
from dotenv import load_dotenv
from tqdm import tqdm
from typing import Optional

load_dotenv()
app = typer.Typer()


@app.command()
def download(
    url: str,
    filename: Optional[str] = typer.Option(
        None, "--filename", "-f", help="Output filename"
    ),
):
    """Download model with correct filename via CivitAI API."""
    # 1. Auth
    key = os.getenv("CIVITAI_API_KEY")
    if not key:
        key = typer.prompt("Enter Civitai API Key", hide_input=True)
        with open(".env", "a") as f:
            f.write(f"\nCIVITAI_API_KEY={key}")

    headers = {"Authorization": f"Bearer {key}"}

    # 2. Identify ID
    # If URL has "modelVersionId=123", use that. Else if "models/123", use that.
    vid_match = re.search(r"modelVersionId=(\d+)", url)
    mid_match = re.search(r"models/(\d+)", url)

    try:
        if vid_match:
            # Get specific version metadata
            rep = requests.get(
                f"https://civitai.com/api/v1/model-versions/{vid_match.group(1)}",
                headers=headers,
            )
            rep.raise_for_status()
            version_data = rep.json()
        elif mid_match:
            # Get model metadata -> extract latest version
            print("Finding latest version...", file=sys.stderr)
            rep = requests.get(
                f"https://civitai.com/api/v1/models/{mid_match.group(1)}",
                headers=headers,
            )
            rep.raise_for_status()
            version_data = rep.json()["modelVersions"][0]
        else:
            print("Error: Invalid URL", file=sys.stderr)
            return

        # 3. Pick the best file (primary or first available)
        files = version_data.get("files", [])
        target = next(
            (f for f in files if f.get("primary")), files[0] if files else None
        )

        if not target:
            print("Error: No files found for this model.", file=sys.stderr)
            return

        download_filename = target["name"]
        download_url = target["downloadUrl"]

        name_part, ext_part = os.path.splitext(download_filename)

        # 4. Determine Filename
        if filename is None and sys.stdout.isatty():
            # Interactive: Prompt with just the name part as default
            filename = typer.prompt(
                "Output filename", default=name_part, show_default=True
            )
        elif filename is None:
            # Pipe: Use full original name
            filename = download_filename

        # Auto-append extension if missing (unless piping)
        if sys.stdout.isatty() and not filename.endswith(ext_part):
            filename += ext_part

    except Exception as e:
        print(f"Metadata error: {e}", file=sys.stderr)
        return

    print("Downloading...", file=sys.stderr)

    with requests.get(download_url, headers=headers, stream=True) as rep:
        rep.raise_for_status()
        total = int(rep.headers.get("Content-Length", 0))

        with tqdm(total=total, unit="B", unit_scale=True, file=sys.stderr) as bar:
            if not sys.stdout.isatty():
                # Pipe mode: Write bytes to stdout, update bar manually
                for chunk in rep.iter_content(8192):
                    sys.stdout.buffer.write(chunk)
                    bar.update(len(chunk))
            else:
                # File mode: Write bytes to file, update bar manually
                with open(filename, "wb") as f:
                    for chunk in rep.iter_content(8192):
                        f.write(chunk)
                        bar.update(len(chunk))


if __name__ == "__main__":
    app()