46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
import mlx.core as mx
|
|
from mlx import nn
|
|
|
|
from mflux.config.runtime_config import RuntimeConfig
|
|
from mflux.post_processing.array_util import ArrayUtil
|
|
from mflux.post_processing.image_util import ImageUtil
|
|
|
|
|
|
class LatentCreator:
|
|
@staticmethod
|
|
def create(
|
|
seed: int,
|
|
height: int,
|
|
width: int,
|
|
):
|
|
return mx.random.normal(
|
|
shape=[1, (height // 16) * (width // 16), 64],
|
|
key=mx.random.key(seed)
|
|
) # fmt: off
|
|
|
|
@staticmethod
|
|
def create_for_txt2img_or_img2img(
|
|
seed: int,
|
|
runtime_conf: RuntimeConfig,
|
|
vae: nn.Module,
|
|
):
|
|
noise = LatentCreator.create(
|
|
seed=seed,
|
|
height=runtime_conf.height,
|
|
width=runtime_conf.width,
|
|
)
|
|
|
|
if runtime_conf.config.init_image_path is None:
|
|
# Text2Image
|
|
return noise
|
|
else:
|
|
# Image2Image
|
|
user_image = ImageUtil.load_image(runtime_conf.config.init_image_path).convert("RGB")
|
|
scaled_user_image = ImageUtil.scale_to_dimensions(user_image, runtime_conf.width, runtime_conf.height)
|
|
encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
|
|
latents = ArrayUtil.pack_latents(encoded, runtime_conf.width, runtime_conf.height)
|
|
sigmas_for_init_image_strength = runtime_conf.sigmas[runtime_conf.init_time_step]
|
|
latents_adjusted = latents * (1.0 - sigmas_for_init_image_strength)
|
|
noise_adjusted = noise * sigmas_for_init_image_strength
|
|
return latents_adjusted + noise_adjusted
|