50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
from dataclasses import dataclass
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import logging
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def get_sigmas(num_inference_steps):
|
|
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
|
|
sigmas = mx.array(sigmas).astype(mx.float32)
|
|
return mx.concatenate([sigmas, mx.zeros(1)])
|
|
|
|
def shift_sigmas(sigmas, width, height):
|
|
y1 = 0.5
|
|
x1 = 256
|
|
m = (1.15 - y1) / (4096 - x1)
|
|
b = y1 - m * x1
|
|
mu = m * width * height / 256 + b
|
|
shifted_sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / sigmas - 1))
|
|
shifted_sigmas[-1] = 0
|
|
return shifted_sigmas
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
precision: mx.Dtype = mx.bfloat16
|
|
|
|
def __init__(
|
|
self,
|
|
num_train_steps: int = 1000,
|
|
num_inference_steps: int = 4,
|
|
width: int = 1024,
|
|
height: int = 1024,
|
|
guidance: float = 4.0,
|
|
):
|
|
self.num_train_steps = num_train_steps
|
|
if width % 16 != 0 or height % 16 != 0:
|
|
log.warning("Width and height should be multiples of 16. Rounding down.")
|
|
self.width = 16 * (height // 16)
|
|
self.height = 16 * (width // 16)
|
|
self.num_inference_steps = num_inference_steps
|
|
self.guidance = guidance
|
|
|
|
def __post_init__(self, **data):
|
|
super().__init__(**data)
|
|
self.__config__.frozen = True
|
|
|
|
|