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

Add stepwise handler

Browse Source
This commit is contained in:
filipstrand 2024-10-09 16:48:59 +02:00
parent f4d0bc799d
commit e908c1cede
9 changed files with 123 additions and 96 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/mflux/__init__.py
View File

@ -3,7 +3,7 @@ from mflux.config.config import ConfigControlnet
from mflux.config.model_config import ModelConfig
from mflux.controlnet.flux_controlnet import Flux1Controlnet
from mflux.flux.flux import Flux1
from mflux.exceptions import StopImageGenerationException
from mflux.error.exceptions import StopImageGenerationException
from mflux.post_processing.image_util import ImageUtil
__all__ = [

71
src/mflux/controlnet/flux_controlnet.py
View File

@ -10,12 +10,14 @@ from mflux.config.runtime_config import RuntimeConfig
from mflux.controlnet.controlnet_util import ControlnetUtil
from mflux.controlnet.transformer_controlnet import TransformerControlnet
from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
from mflux.exceptions import StopImageGenerationException
from mflux.error.exceptions import StopImageGenerationException
from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from mflux.models.transformer.transformer import Transformer
from mflux.models.vae.vae import VAE
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.image_util import ImageUtil
from mflux.post_processing.stepwise_handler import StepwiseHandler
from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
from mflux.tokenizer.t5_tokenizer import TokenizerT5
from mflux.tokenizer.tokenizer_handler import TokenizerHandler
@ -112,11 +114,18 @@ class Flux1Controlnet:
controlnet_save_canny: bool = False,
config: ConfigControlnet = ConfigControlnet(),
stepwise_output_dir: Path = None,
stepwise_composite_only=False
) -> "GeneratedImage": # fmt: off
# Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.num_inference_steps))
stepwise_handler = StepwiseHandler(
flux=self,
config=config,
seed=seed,
prompt=prompt,
time_steps=time_steps,
output_dir=stepwise_output_dir,
)
# Embedd the controlnet reference image
control_image = ImageUtil.load_image(controlnet_image_path)
@ -127,7 +136,7 @@ class Flux1Controlnet:
controlnet_cond = ImageUtil.to_array(control_image)
controlnet_cond = self.vae.encode(controlnet_cond)
controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
controlnet_cond = ArrayUtil.pack_latents(controlnet_cond, config.height, config.width)
# 1. Create the initial latents
latents = mx.random.normal(
@ -135,15 +144,12 @@ class Flux1Controlnet:
key=mx.random.key(seed)
) # fmt: off
# 2. Embedd the prompt
# 2. Embed the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)
clip_tokens = self.clip_tokenizer.tokenize(prompt)
prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
step_wise_images = []
if stepwise_output_dir:
stepwise_output_dir.mkdir(parents=True, exist_ok=True)
for t in time_steps:
try:
# Compute controlnet samples
@ -171,41 +177,18 @@ class Flux1Controlnet:
dt = config.sigmas[t + 1] - config.sigmas[t]
latents += noise * dt
# Handle stepwise output if enabled
stepwise_handler.process_step(t, latents)
# Evaluate to enable progress tracking
mx.eval(latents)
if stepwise_output_dir:
stepwise_decoded = self.vae.decode(
Flux1Controlnet._unpack_latents(latents, config.height, config.width)
)
# performance todo: Pillow mostly uses CPU,
# can try to improve image generation performance by offloading
# stepwise image processing to the CPU via threading
stepwise_img = ImageUtil.to_image(
decoded_latents=stepwise_decoded,
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,
)
step_wise_images.append(stepwise_img)
if not stepwise_composite_only:
stepwise_img.save(
path=stepwise_output_dir / f"seed_{seed}_step{t+1}of{len(time_steps)}.png",
export_json_metadata=False,
)
except KeyboardInterrupt: # noqa: PERF203
raise StopImageGenerationException(f"Stopping image generation at step {t+1}/{len(time_steps)}")
finally:
if step_wise_images and stepwise_output_dir:
composite_img = ImageUtil.to_composite_image(step_wise_images)
composite_img.save(stepwise_output_dir / f"seed_{seed}_composite.png")
except KeyboardInterrupt:
stepwise_handler.handle_interruption()
raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
# 5. Decode the latent array and return the image
latents = Flux1Controlnet._unpack_latents(latents, config.height, config.width)
latents = ArrayUtil.unpack_latents(latents, config.height, config.width)
decoded = self.vae.decode(latents)
return ImageUtil.to_image(
decoded_latents=decoded,
@ -219,20 +202,6 @@ class Flux1Controlnet:
controlnet_image_path=controlnet_image_path,
)
@staticmethod
def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
return latents
@staticmethod
def _pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))
latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5))
latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64))
return latents
def _set_model_weights(self, weights):
self.vae.update(weights.vae)
self.transformer.update(weights.transformer)

0
src/mflux/error/__init__.py Normal file
View File

0
src/mflux/exceptions.py → src/mflux/error/exceptions.py
View File

70
src/mflux/flux/flux.py
View File

@ -1,16 +1,19 @@
import mlx.core as mx
from pathlib import Path
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.config.runtime_config import RuntimeConfig
from mflux.exceptions import StopImageGenerationException
from mflux.error.exceptions import StopImageGenerationException
from mflux.post_processing.stepwise_handler import StepwiseHandler
from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from mflux.models.transformer.transformer import Transformer
from mflux.models.vae.vae import VAE
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.generated_image import GeneratedImage
from mflux.post_processing.image_util import ImageUtil
from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
@ -71,10 +74,24 @@ class Flux1:
if weights.quantization_level is not None:
self._set_model_weights(weights)
def generate_image(self, seed: int, prompt: str, config: Config = Config(), stepwise_output_dir: Path = None, stepwise_composite_only=False) -> GeneratedImage: # fmt: off
def generate_image(
self,
seed: int,
prompt: str,
config: Config = Config(),
stepwise_output_dir: Path = None,
) -> GeneratedImage:
# Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.num_inference_steps))
stepwise_handler = StepwiseHandler(
flux=self,
config=config,
seed=seed,
prompt=prompt,
time_steps=time_steps,
output_dir=stepwise_output_dir,
)
# 1. Create the initial latents
latents = mx.random.normal(
@ -82,15 +99,12 @@ class Flux1:
key=mx.random.key(seed)
) # fmt: off
# 2. Embedd the prompt
# 2. Embed the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)
clip_tokens = self.clip_tokenizer.tokenize(prompt)
prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
step_wise_images = []
if stepwise_output_dir:
stepwise_output_dir.mkdir(parents=True, exist_ok=True)
for t in time_steps:
try:
# 3.t Predict the noise
@ -106,39 +120,18 @@ class Flux1:
dt = config.sigmas[t + 1] - config.sigmas[t]
latents += noise * dt
# Handle stepwise output if enabled
stepwise_handler.process_step(t, latents)
# Evaluate to enable progress tracking
mx.eval(latents)
if stepwise_output_dir:
stepwise_decoded = self.vae.decode(Flux1._unpack_latents(latents, config.height, config.width))
# performance todo: Pillow mostly uses CPU,
# can try to improve image generation performance by offloading
# stepwise image processing to the CPU via threading
stepwise_img = ImageUtil.to_image(
decoded_latents=stepwise_decoded,
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,
)
step_wise_images.append(stepwise_img)
if not stepwise_composite_only:
stepwise_img.save(
path=stepwise_output_dir / f"seed_{seed}_step{t+1}of{len(time_steps)}.png",
export_json_metadata=False,
)
except KeyboardInterrupt: # noqa: PERF203
raise StopImageGenerationException(f"Stopping image generation at step {t+1}/{len(time_steps)}")
finally:
if step_wise_images and stepwise_output_dir:
composite_img = ImageUtil.to_composite_image(step_wise_images)
composite_img.save(stepwise_output_dir / f"seed_{seed}_composite.png")
except KeyboardInterrupt:
stepwise_handler.handle_interruption()
raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
# 5. Decode the latent array and return the image
latents = Flux1._unpack_latents(latents, config.height, config.width)
latents = ArrayUtil.unpack_latents(latents, config.height, config.width)
decoded = self.vae.decode(latents)
return ImageUtil.to_image(
decoded_latents=decoded,
@ -151,13 +144,6 @@ class Flux1:
config=config,
)
@staticmethod
def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
return latents
@staticmethod
def from_alias(alias: str, quantize: int | None = None) -> "Flux1":
return Flux1(

2
src/mflux/generate.py
View File

@ -46,13 +46,13 @@ def main():
image = flux.generate_image(
seed=int(time.time()) if args.seed is None else args.seed,
prompt=args.prompt,
stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
config=Config(
num_inference_steps=args.steps,
height=args.height,
width=args.width,
guidance=args.guidance,
),
stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
)
# Save the image

2
src/mflux/generate_controlnet.py
View File

@ -52,6 +52,7 @@ def main():
output=args.output,
controlnet_image_path=args.controlnet_image_path,
controlnet_save_canny=args.controlnet_save_canny,
stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
config=ConfigControlnet(
num_inference_steps=args.steps,
height=args.height,
@ -59,7 +60,6 @@ def main():
guidance=args.guidance,
controlnet_strength=args.controlnet_strength,
),
stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
)
# Save the image

17
src/mflux/post_processing/array_util.py Normal file
View File

@ -0,0 +1,17 @@
import mlx.core as mx
class ArrayUtil:
@staticmethod
def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
return latents
@staticmethod
def pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))
latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5))
latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64))
return latents

55
src/mflux/post_processing/stepwise_handler.py Normal file
View File

@ -0,0 +1,55 @@
from pathlib import Path
import mlx.core as mx
from mflux.config.runtime_config import RuntimeConfig
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.image_util import ImageUtil
class StepwiseHandler:
def __init__(
self,
flux,
config: RuntimeConfig,
seed: int,
prompt: str,
time_steps,
output_dir: Path | None = None,
):
self.flux = flux
self.config = config
self.seed = seed
self.prompt = prompt
self.output_dir = output_dir
self.time_steps = time_steps
self.step_wise_images = []
if self.output_dir:
self.output_dir.mkdir(parents=True, exist_ok=True)
def process_step(self, step: int, latents: mx.array):
if self.output_dir:
unpack_latents = ArrayUtil.unpack_latents(latents, self.config.height, self.config.width)
stepwise_decoded = self.flux.vae.decode(unpack_latents)
stepwise_img = ImageUtil.to_image(
decoded_latents=stepwise_decoded,
seed=self.seed,
prompt=self.prompt,
quantization=self.flux.bits,
generation_time=self.time_steps.format_dict["elapsed"],
lora_paths=self.flux.lora_paths,
lora_scales=self.flux.lora_scales,
config=self.config,
)
self.step_wise_images.append(stepwise_img)
stepwise_img.save(
path=self.output_dir / f"seed_{self.seed}_step{step + 1}of{len(self.time_steps)}.png",
export_json_metadata=False,
)
def handle_interruption(self):
if self.step_wise_images:
composite_img = ImageUtil.to_composite_image(self.step_wise_images)
composite_img.save(self.output_dir / f"seed_{self.seed}_composite.png")