from enum import Enum

from pydantic import BaseModel


class HDStrategy(str, Enum):
    ORIGINAL = 'Original'
    RESIZE = 'Resize'
    CROP = 'Crop'


class LDMSampler(str, Enum):
    ddim = 'ddim'
    plms = 'plms'


class SDSampler(str, Enum):
    ddim = 'ddim'
    pndm = 'pndm'


class Config(BaseModel):
    ldm_steps: int
    ldm_sampler: str = LDMSampler.plms
    zits_wireframe: bool = True
    hd_strategy: str
    hd_strategy_crop_margin: int
    hd_strategy_crop_trigger_size: int
    hd_strategy_resize_limit: int

    prompt: str = ''
    # 始终是在原图尺度上的值
    use_croper: bool = False
    croper_x: int = None
    croper_y: int = None
    croper_height: int = None
    croper_width: int = None

    # sd
    sd_strength: float = 0.75
    sd_steps: int = 50
    sd_guidance_scale: float = 7.5
    sd_sampler: str = SDSampler.ddim
    # -1 mean random seed
    sd_seed: int = 42