33 lines
735 B
Python
33 lines
735 B
Python
import ml_collections
|
|
import imp
|
|
import os
|
|
|
|
base = imp.load_source("base", os.path.join(os.path.dirname(__file__), "base.py"))
|
|
|
|
|
|
def get_config():
|
|
config = base.get_config()
|
|
|
|
config.pretrained.model = "runwayml/stable-diffusion-v1-5"
|
|
|
|
config.mixed_precision = "fp16"
|
|
config.allow_tf32 = True
|
|
config.use_lora = False
|
|
|
|
config.train.batch_size = 4
|
|
config.train.gradient_accumulation_steps = 2
|
|
config.train.learning_rate = 3e-5
|
|
config.train.clip_range = 1e-4
|
|
|
|
# sampling
|
|
config.sample.num_steps = 50
|
|
config.sample.batch_size = 8
|
|
config.sample.num_batches_per_epoch = 4
|
|
|
|
config.per_prompt_stat_tracking = {
|
|
"buffer_size": 16,
|
|
"min_count": 16,
|
|
}
|
|
|
|
return config
|