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

Rewrite img2img latent creation to emphasize the linear interp between latent and noise

Browse Source
This commit is contained in:
filipstrand 2024-10-29 19:19:59 +01:00
parent 8ac08f7853
commit 4e2a10b1b9
1 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/mflux/latent_creator/latent_creator.py
View File

@ -12,7 +12,7 @@ class LatentCreator:
seed: int, seed: int,
height: int, height: int,
width: int, width: int,
): ) -> mx.array:
return mx.random.normal( return mx.random.normal(
shape=[1, (height // 16) * (width // 16), 64], shape=[1, (height // 16) * (width // 16), 64],
key=mx.random.key(seed) key=mx.random.key(seed)
@ -23,8 +23,8 @@ class LatentCreator:
seed: int, seed: int,
runtime_conf: RuntimeConfig, runtime_conf: RuntimeConfig,
vae: nn.Module, vae: nn.Module,
): ) -> mx.array:
noise = LatentCreator.create( pure_noise = LatentCreator.create(
seed=seed, seed=seed,
height=runtime_conf.height, height=runtime_conf.height,
width=runtime_conf.width, width=runtime_conf.width,
@ -32,14 +32,20 @@ class LatentCreator:
if runtime_conf.config.init_image_path is None: if runtime_conf.config.init_image_path is None:
# Text2Image # Text2Image
return noise return pure_noise
else: else:
# Image2Image # Image2Image
user_image = ImageUtil.load_image(runtime_conf.config.init_image_path).convert("RGB") 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) scaled_user_image = ImageUtil.scale_to_dimensions(user_image, runtime_conf.width, runtime_conf.height)
encoded = vae.encode(ImageUtil.to_array(scaled_user_image)) encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
latents = ArrayUtil.pack_latents(encoded, runtime_conf.width, runtime_conf.height) latents = ArrayUtil.pack_latents(encoded, runtime_conf.width, runtime_conf.height)
sigmas_for_init_image_strength = runtime_conf.sigmas[runtime_conf.init_time_step] sigma = runtime_conf.sigmas[runtime_conf.init_time_step]
latents_adjusted = latents * (1.0 - sigmas_for_init_image_strength) return LatentCreator.add_noise_by_interpolation(
noise_adjusted = noise * sigmas_for_init_image_strength clean=latents,
return latents_adjusted + noise_adjusted noise=pure_noise,
sigma=sigma
) # fmt: off
@staticmethod
def add_noise_by_interpolation(clean: mx.array, noise: mx.array, sigma: float) -> mx.array:
return (1 - sigma) * clean + sigma * noise