cmhub/apps/ai/management/commands/import_ai_models.py

import json
from pathlib import Path

from django.core.management.base import BaseCommand, CommandError

from apps.ai.importers import import_ai_models_config


class Command(BaseCommand):
    help = "Import cmbot-style ai_models.json into encrypted AiModel records."

    def add_arguments(self, parser):
        parser.add_argument("path", help="Path to ai_models.json.")
        parser.add_argument(
            "--create-default-aliases",
            action="store_true",
            help="Create title-standard and image-hd default aliases when possible.",
        )

    def handle(self, *args, **options):
        path = Path(options["path"])
        if not path.exists():
            raise CommandError(f"File not found: {path}")

        try:
            config = json.loads(path.read_text(encoding="utf-8-sig"))
            result = import_ai_models_config(
                config,
                create_default_aliases=options["create_default_aliases"],
            )
        except Exception as exc:
            raise CommandError(str(exc)) from exc

        self.stdout.write(
            self.style.SUCCESS(
                "Imported AI models: "
                f"created={result['created']}, "
                f"updated={result['updated']}, "
                f"aliases={result['aliases']}"
            )
        )