from airunner.settings import AIRUNNER_ART_ENABLED


ai_art_models = [
    {
        "name": "Stable Diffusion 1.5",
        "path": "stable-diffusion-v1-5/stable-diffusion-v1-5",
        "branch": "main",
        "version": "SD 1.5",
        "category": "stablediffusion",
        "pipeline_action": "txt2img",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "Stable Diffusion XL Base 1.0",
        "path": "stabilityai/stable-diffusion-xl-base-1.0",
        "branch": "main",
        "version": "SDXL 1.0",
        "category": "stablediffusion",
        "pipeline_action": "txt2img",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "Stable Diffusion XL Turbo",
        "path": "stabilityai/stable-diffusion-xl-base-1.0",
        "branch": "main",
        "version": "SDXL Turbo",
        "category": "stablediffusion",
        "pipeline_action": "txt2img",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "Stable Diffusion Inpaint 1.5",
        "path": "stable-diffusion-v1-5/stable-diffusion-inpainting",
        "branch": "fp16",
        "version": "SD 1.5",
        "category": "stablediffusion",
        "pipeline_action": "inpaint",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "SDXL Inpaint",
        "path": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
        "branch": "fp16",
        "version": "SDXL 1.0",
        "category": "stablediffusion",
        "pipeline_action": "inpaint",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "CompVis Safety Checker",
        "path": "CompVis/stable-diffusion-safety-checker",
        "branch": "main",
        "version": "SD 1.5",
        "category": "stablediffusion",
        "pipeline_action": "txt2img",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "OpenAI Feature Extractor",
        "path": "openai/clip-vit-large-patch14",
        "branch": "main",
        "version": "SD 1.5",
        "category": "stablediffusion",
        "pipeline_action": "txt2img",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "Inpaint vae",
        "path": "cross-attention/asymmetric-autoencoder-kl-x-1-5",
        "branch": "fp16",
        "version": "SD 1.5",
        "category": "stablediffusion",
        "pipeline_action": "inpaint_vae",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
    {
        "name": "Flux S",
        "path": "black-forest-labs/FLUX.1-schnell",
        "branch": "main",
        "version": "Flux S",
        "category": "flux",
        "pipeline_action": "txt2img",
        "enabled": True,
        "model_type": "art",
        "is_default": True,
    },
]

llm_models = [
    {
        "name": "Mistral 7B Instruct v3",
        "path": "w4ffl35/Ministral-8B-Instruct-2410-doublequant",
        "branch": "main",
        "version": "0.3",
        "category": "llm",
        "pipeline_action": "causallm",
        "enabled": True,
        "model_type": "llm",
        "is_default": True,
    },
    {
        "name": "Intfloat E5 Large",
        "path": "intfloat/e5-large",
        "branch": "main",
        "version": "llm",
        "category": "llm",
        "pipeline_action": "embedding",
        "enabled": True,
        "model_type": "llm",
        "is_default": True,
    },
]

if AIRUNNER_ART_ENABLED:
    model_bootstrap_data = ai_art_models + llm_models
else:
    model_bootstrap_data = llm_models
