from dataclasses import dataclass
from typing import Optional

DEFAULT_MODEL = "SD-1.5"
DEFAULT_SAMPLER = "k_dpmpp_2m"

DEFAULT_NEGATIVE_PROMPT = (
    "(weird eyes:1.5), (smudged face:1.8), (blurred face:1.5), (poorly drawn face:1.5), mutation, mutilation, (cloned face:2.0), (strange mouth:1.5), (duplicated face:2.0), "
    "(Ugly:1.5), (duplication:1.5), duplicates, mutilation, deformed, mutilated, (mutation:1.5), (twisted body:1.5), (disfigured:1.5), (bad anatomy:1.5), "
    "(poorly drawn hands:1.5), (extra limbs:1.5), (malformed limbs:1.5), (missing arms:1.5), (extra arms:1.5), (missing legs:1.5), (extra legs:1.5), (mutated hands:1.5), "
    "(extra hands:1.5), (fused fingers:1.5), (missing fingers:1.5), extra fingers, (long neck:1.8), small head, (closed eyes:1.5), rolling eyes, "
    "grainy, blurred, blurry, (writing:1.5), calligraphy, (text:3.0), (watermark:1.5), bad art, missing fingers, (missing limbs:1.5), "
    "(out of frame:1.5), (extra fingers:1.5), (mutated hands:1.5), (porn:1.5), (child abuse:1.5), (horrific images:1.5), (dead bodies:1.5), "
    "(watermark:1.5), child, (lowres:1.5), bad anatomy, bad hands, text, error, (missing fingers:1.5), (extra digit:1.5), (fewer digits:1.6), (cropped:1.5), "
    "worst quality, (low quality:1.5), normal quality, jpeg artifacts, (signature:1.5), watermark, username, blurry, (bad feet:1.5), (bad lighting:1.3), "
    "bad coloring, bad perspective, (bad focus:1.5), (bad design:1.3), (bad eyes:1.7), (bad hair:1.7), bad nose, bad mouth, bad ears, bad teeth, bad tongue, bad lips, "
    "(nude:1.5), (naked:1.5), (paedophile:2.0), (pedo:2.0), (floating limbs: 1.5), (bad art, low detail:1.5), (plain background, grainy, low "
    "quality, mutated hands and fingers:1.4), (watermark, thin lines:1.2), (deformed, signature:1.2), (blurry, ugly, bad anatomy, extra"
    "limbs, undersaturated, low resolution:1.5), (disfigured, deformations, out of frame:1.5), (amputee:1.8), (bad proportions:1.5), (extra limb:1.5), (missing"
    "limbs:1.4), (distortion:1.5), (out of frame:2.0), (poorly drawn face:2.0), (poorly drawn hands:1.5), (text:1.8), malformed, (cropped:1.5)"
)

SPLITMEM_ENABLED = False


@dataclass
class ModelConfig:
    short_name: str
    config_path: str
    weights_url: str
    default_image_size: int
    weights_url_full: str = None
    forced_attn_precision: str = "default"
    default_negative_prompt: str = DEFAULT_NEGATIVE_PROMPT
    alias: str = None


midas_url = "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt"

MODEL_CONFIGS = [
    ModelConfig(
        short_name="SD-1.4",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/bstddev/sd-v1-4/resolve/77221977fa8de8ab8f36fac0374c120bd5b53287/sd-v1-4.ckpt",
        default_image_size=512,
        alias="sd14",
    ),
    ModelConfig(
        short_name="SD-1.5",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/889b629140e71758e1e0006e355c331a5744b4bf/v1-5-pruned-emaonly.ckpt",
        weights_url_full="https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/889b629140e71758e1e0006e355c331a5744b4bf/v1-5-pruned.ckpt",
        default_image_size=512,
        alias="sd15",
    ),
    ModelConfig(
        short_name="SD-1.5-inpaint",
        config_path="configs/stable-diffusion-v1-inpaint.yaml",
        weights_url="https://huggingface.co/julienacquaviva/inpainting/resolve/2155ff7fe38b55f4c0d99c2f1ab9b561f8311ca7/sd-v1-5-inpainting.ckpt",
        default_image_size=512,
        alias="sd15in",
    ),
    ModelConfig(
        short_name="SD-2.0",
        config_path="configs/stable-diffusion-v2-inference.yaml",
        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt",
        default_image_size=512,
        alias="sd20",
    ),
    ModelConfig(
        short_name="SD-2.0-inpaint",
        config_path="configs/stable-diffusion-v2-inpainting-inference.yaml",
        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.ckpt",
        default_image_size=512,
        alias="sd20in",
    ),
    ModelConfig(
        short_name="SD-2.1",
        config_path="configs/stable-diffusion-v2-inference.yaml",
        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt",
        default_image_size=512,
        alias="sd21",
    ),
    ModelConfig(
        short_name="SD-2.1-inpaint",
        config_path="configs/stable-diffusion-v2-inpainting-inference.yaml",
        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.ckpt",
        default_image_size=512,
        alias="sd21in",
    ),
    ModelConfig(
        short_name="SD-2.1-v",
        config_path="configs/stable-diffusion-v2-inference-v.yaml",
        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.ckpt",
        default_image_size=768,
        forced_attn_precision="fp32",
    ),
    ModelConfig(
        short_name="SD-2.0-v",
        config_path="configs/stable-diffusion-v2-inference-v.yaml",
        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2/resolve/main/768-v-ema.ckpt",
        default_image_size=768,
        alias="sd20v",
    ),
    ModelConfig(
        short_name="SD-2.0-depth",
        config_path="configs/stable-diffusion-v2-midas-inference.yaml",
        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-depth/resolve/main/512-depth-ema.ckpt",
        default_image_size=512,
        alias="sd20dep",
    ),
    ModelConfig(
        short_name="instruct-pix2pix",
        config_path="configs/instruct-pix2pix.yaml",
        weights_url="https://huggingface.co/imaginairy/instruct-pix2pix/resolve/ea0009b3d0d4888f410a40bd06d69516d0b5a577/instruct-pix2pix-00-22000-pruned.ckpt",
        default_image_size=512,
        default_negative_prompt="",
        alias="edit",
    ),
    # --- Custom model --#
    ModelConfig(
        short_name="openjourney-v1",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/prompthero/openjourney/resolve/7428477dad893424c92f6ea1cc29d45f6d1448c1/mdjrny-v4.safetensors",
        default_image_size=512,
        default_negative_prompt="",
        alias="oj1",
    ),
    ModelConfig(
        short_name="openjourney-v2",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/prompthero/openjourney-v2/resolve/47257274a40e93dab7fbc0cd2cfd5f5704cfeb60/openjourney-v2.ckpt",
        default_image_size=512,
        default_negative_prompt="",
        alias="oj2",
    ),
    ModelConfig(
        short_name="MidJourney-PaperCut",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/ShadoWxShinigamI/MidJourney-PaperCut/resolve/main/Mdjrny-pprct_step_7000.ckpt",
        default_image_size=512,
        alias="mjpc"
    ),
    ModelConfig(
        short_name="Openjourney",
        config_path="configs/stable-diffusion-v1.yaml",
        # config_path="configs/stable-diffusion-v2-inference-v.yaml",
        # weights_url="https://civitai.com/api/download/models/96",
        weights_url="https://huggingface.co/prompthero/openjourney-v2/resolve/main/openjourney-v2-unpruned.ckpt",
        default_image_size=512,
        alias="ojc"
    ),
    ModelConfig(
        short_name="Protogen-Anime",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://civitai.com/api/download/models/4007",
        default_image_size=512,
        alias="proa"
    ),
    ModelConfig(
        short_name="Protogen-Photo",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/darkstorm2150/Protogen_x3.4_Official_Release/resolve/main/ProtoGen_X3.4.ckpt",
        default_image_size=512,
        alias="prop"
    ),
    ModelConfig(
        short_name="Protogen-PhotoV2",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/darkstorm2150/Protogen_v5.3_Official_Release/resolve/main/ProtoGen_X5.3.ckpt",
        default_image_size=512,
        alias="prop2"
    ),
    ModelConfig(
        short_name="RPG",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/Anashel/rpg/resolve/main/RPG-V4-Model-Download/RPG-v4.safetensors",
        default_image_size=512,
        alias="rpg"
    ),
    ModelConfig(
        short_name="SynthwavePunk",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/zipp425/synthwavePunk/resolve/main/synthwavePunk_v2.ckpt",
        default_image_size=512,
    ),
    ModelConfig(
        short_name="ModernDisney",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/nitrosocke/mo-di-diffusion/resolve/main/moDi-v1-pruned.ckpt",
        default_image_size=512,
    ),
    ModelConfig(
        short_name="Art-Diffusion",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/DGSpitzer/DGSpitzer-Art-Diffusion/resolve/main/DGSpitzer-Art-Diffusion.ckpt",
        default_image_size=512,
    ),
    ModelConfig(
        short_name="graphic-art",
        config_path="configs/graphicArt_graphicArtBeta.yaml",
        weights_url="https://civitai.com/api/download/models/9297?type=Model&format=PickleTensor",
        default_image_size=768,
    ),
    ModelConfig(
        short_name="DreamLikeHumans",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/dreamlike-art/dreamlike-photoreal-2.0/resolve/main/dreamlike-photoreal-2.0.ckpt",
        default_image_size=512,
        alias="dlh"
    ),
    ModelConfig(
        short_name="DreamLikeDiffusion",
        config_path="configs/stable-diffusion-v1.yaml",
        weights_url="https://huggingface.co/dreamlike-art/dreamlike-diffusion-1.0/resolve/main/dreamlike-diffusion-1.0.ckpt",
        default_image_size=640,
        alias="dld"
    ),
    # ModelConfig(
    #     short_name="SD-2.0-upscale",
    #     config_path="configs/stable-diffusion-v2-upscaling.yaml",
    #     weights_url="https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler/resolve/main/x4-upscaler-ema.ckpt",
    #     default_image_size=512,
    # ),
]

MODEL_CONFIG_SHORTCUTS = {m.short_name: m for m in MODEL_CONFIGS}
for m in MODEL_CONFIGS:
    if m.alias:
        MODEL_CONFIG_SHORTCUTS[m.alias] = m

MODEL_CONFIG_SHORTCUTS["openjourney"] = MODEL_CONFIG_SHORTCUTS["openjourney-v2"]
MODEL_CONFIG_SHORTCUTS["oj"] = MODEL_CONFIG_SHORTCUTS["openjourney-v2"]

MODEL_SHORT_NAMES = sorted(MODEL_CONFIG_SHORTCUTS.keys())

SAMPLER_TYPE_OPTIONS = [
    "plms",
    "ddim",
    "k_dpm_fast",
    "k_dpm_adaptive",
    "k_lms",
    "k_dpm_2",
    "k_dpm_2_a",
    "k_dpmpp_2m",
    "k_dpmpp_2s_a",
    "k_euler",
    "k_euler_a",
    "k_heun",
]


def get_model_config(model_name: str) -> Optional[ModelConfig]:
    if model_name in MODEL_CONFIG_SHORTCUTS:
        return MODEL_CONFIG_SHORTCUTS[model_name]
    return None
