ai-art

  1import os
  2
  3from clip.clip import available_models
  4
  5from typing import List
  6from dataclasses import dataclass, field
  7
  8
  9INIT_NOISES = ["", "gradient", "pixels", "fractal"]
 10OPTIMIZERS = ["Adam", "AdamW", "Adagrad", "Adamax"]
 11AUGMENTS = ["Ji", "Sh", "Gn", "Pe", "Ro", "Af", "Et", "Ts", "Cr", "Er", "Re", "Hf"]
 12
 13
 14@dataclass
 15class Config:
 16    prompts: List[str] = field(default_factory=lambda: [])
 17    image_prompts: List[str] = field(default_factory=lambda: [])
 18    max_iterations: int = 500
 19    save_freq: int = 50
 20    size: List[int] = field(default_factory=lambda: [256, 256])
 21    pixelart: List[int] = None
 22    init_image: str = ""
 23    init_noise: str = "gradient"
 24    init_weight: float = 0.0
 25    mse_decay_rate: float = 0.0
 26    output_dir: str = "/data/outputs"
 27    models_dir: str = "/data/models"
 28    clip_model: str = "ViT-B/16"
 29    vqgan_checkpoint: str = "/data/models/vqgan_imagenet_f16_16384.ckpt"
 30    vqgan_config: str = "/data/models/vqgan_imagenet_f16_16384.json"
 31    noise_prompt_seeds: List[int] = field(default_factory=lambda: [])
 32    noise_prompt_weights: List[float] = field(default_factory=lambda: [])
 33    step_size: float = 0.1
 34    cutn: int = 32
 35    cut_pow: float = 1.0
 36    seed: int = -1
 37    optimizer: str = "Adam"
 38    nwarm_restarts: int = -1
 39    augments: List[str] = field(default_factory=lambda: ["Af", "Pe", "Ji", "Er"])
 40
 41    def __post_init__(self):
 42        if self.init_noise not in INIT_NOISES:
 43            exit(
 44                f'ERROR: "init_noise": {self.init_noise}, <-- Noise algorithm not found.\n'
 45                f"Currently only the following values are supported: {INIT_NOISES}."
 46            )
 47
 48        if self.optimizer not in OPTIMIZERS:
 49            exit(
 50                f'ERROR: "optimizer": {self.optimizer}, <-- Optimizer not found.\n'
 51                f"Currently only the following values are supported: {OPTIMIZERS}."
 52            )
 53
 54        os.makedirs(self.models_dir, exist_ok=True)
 55        os.makedirs(self.output_dir, exist_ok=True)
 56        os.makedirs(f"{self.output_dir}/steps", exist_ok=True)
 57        print(f"Saving outputs in '{self.output_dir}'")
 58
 59        models = available_models()
 60        if not os.path.exists(self.clip_model) and self.clip_model not in models:
 61            exit(
 62                f'ERROR: "clip_model": {self.clip_model}, <-- Model not found.\n'
 63                f"Make sure it is a valid path to a downloaded model or match one of {models}."
 64            )
 65
 66        if not os.path.exists(self.vqgan_config):
 67            exit(
 68                f'ERROR: "vqgan_config": {self.vqgan_config}, <-- Configuration file not found.\n'
 69                f"Make sure the path is correct (Multiple config files are available in the `./configs/models` directory)."
 70            )
 71
 72        if not os.path.exists(self.vqgan_checkpoint):
 73            exit(
 74                f'ERROR: "vqgan_checkpoint": {self.vqgan_checkpoint}, <-- Model not found.\n'
 75                f"Make sure the path is correct and that you have downloaded the model (Refer to the README)."
 76            )
 77
 78        if self.pixelart:
 79            print(
 80                "Enabling PixelArt mode. It is recommended to add 'pixelart' to your prompt."
 81            )
 82
 83    def __str__(self):
 84        _str = (
 85            f"Config:\n"
 86            f"  - prompts: {self.prompts}\n"
 87            f"  - image_prompts: {self.image_prompts}\n"
 88            f"  - max_iterations: {self.max_iterations}\n"
 89            f"  - save_freq: {self.save_freq}\n"
 90            f"  - size: {self.size}\n"
 91            f"  - pixelart: {self.pixelart}\n"
 92            f"  - init_image: {self.init_image}\n"
 93            f"  - init_noise: {self.init_noise}\n"
 94            f"  - init_weight: {self.init_weight}\n"
 95            f"  - mse_decay_rate: {self.mse_decay_rate}\n"
 96            f"  - output_dir: {self.output_dir}\n"
 97            f"  - models_dir: {self.models_dir}\n"
 98            f"  - clip_model: {self.clip_model}\n"
 99            f"  - vqgan_checkpoint: {self.vqgan_checkpoint}\n"
100            f"  - vqgan_config: {self.vqgan_config}\n"
101            f"  - noise_prompt_seeds: {self.noise_prompt_seeds}\n"
102            f"  - noise_prompt_weights: {self.noise_prompt_weights}\n"
103            f"  - step_size: {self.step_size}\n"
104            f"  - cutn: {self.cutn}\n"
105            f"  - cut_pow: {self.cut_pow}\n"
106            f"  - seed: {self.seed}\n"
107            f"  - optimizer: {self.optimizer}\n"
108            f"  - nwarm_restarts: {self.nwarm_restarts}\n"
109            f"  - augments: {self.augments}\n"
110        )
111        return _str