monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

make config static and freeze after initialisation

Browse Source
This commit is contained in:
Fabio 2024-08-19 20:41:11 +02:00
parent 40b77e81bc
commit e60cd79f0c
3 changed files with 33 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
src/flux_1_schnell/config/config.py
View File

@ -1,3 +1,4 @@
from dataclasses import dataclass
import mlx.core as mx
import numpy as np
import logging
@ -5,37 +6,44 @@ 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
num_train_steps = 1000
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)
base_sigmas = Config.base_sigmas(num_inference_steps)
self.num_inference_steps = num_inference_steps
self.guidance = guidance
self.sigmas = mx.concatenate([base_sigmas, mx.zeros(1)])
@staticmethod
def base_sigmas(num_inference_steps):
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
sigmas = mx.array(sigmas).astype(mx.float32)
return sigmas
def __post_init__(self, **data):
super().__init__(**data)
self.__config__.frozen = True
def shift_sigmas(self):
y1 = 0.5
x1 = 256
m = (1.15 - y1) / (4096 - x1)
b = y1 - m * x1
mu = m * self.width * self.height / 256 + b
self.sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / self.sigmas - 1))
self.sigmas[-1] = 0

10
src/flux_1_schnell/flux.py
View File

@ -3,7 +3,7 @@ import mlx.core as mx
from PIL import Image
from tqdm import tqdm
from flux_1_schnell.config.config import Config
from flux_1_schnell.config.config import Config, get_sigmas, shift_sigmas
from flux_1_schnell.latent_creator.latent_creator import LatentCreator
from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from flux_1_schnell.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
@ -31,8 +31,9 @@ class Flux1:
self.clip_text_encoder = CLIPEncoder(weights.clip_encoder)
def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image:
sigmas = get_sigmas(config.num_inference_steps)
if self.is_dev:
config.shift_sigmas()
sigmas = shift_sigmas(sigmas)
latents = LatentCreator.create(config.height, config.width, seed)
t5_tokens = self.t5_tokenizer.tokenize(prompt)
@ -46,10 +47,11 @@ class Flux1:
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
hidden_states=latents,
config=config
config=config,
sigmas=sigmas
)
dt = config.sigmas[t + 1] - config.sigmas[t]
dt = sigmas[t + 1] - sigmas[t]
latents += noise * dt
mx.eval(latents)

5
src/flux_1_schnell/models/transformer/transformer.py
View File

@ -32,9 +32,10 @@ class Transformer(nn.Module):
prompt_embeds: mx.array,
pooled_prompt_embeds: mx.array,
hidden_states: mx.array,
config: Config
config: Config,
sigmas: mx.array,
) -> mx.array:
time_step = config.sigmas[t] * config.num_train_steps
time_step = sigmas[t] * config.num_train_steps
time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
hidden_states = self.x_embedder(hidden_states)
guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision)