56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
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")
|