From 5669d26274eb53c9fa263fe28b16faa327928d85 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sun, 18 Aug 2024 20:41:48 +0200 Subject: [PATCH 1/2] Never override previously generated images --- README.md | 3 ++- main.py | 3 ++- src/flux_1_schnell/config/config.py | 1 + .../post_processing/image_util.py | 25 +++++++++++++++++++ 4 files changed, 30 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1d1d1f8..06bebbe 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ sys.path.append("/path/to/mflux/src") from flux_1_schnell.config.config import Config from flux_1_schnell.flux import Flux1Schnell +from flux_1_schnell.post_processing.image_util import ImageUtil flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell") @@ -62,7 +63,7 @@ image = flux.generate_image( ) ) -image.save("image.png") +ImageUtil.save_image(image, "image.png") ``` If the model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell model). diff --git a/main.py b/main.py index ba639a1..5cf39bb 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))) from flux_1_schnell.config.config import Config from flux_1_schnell.flux import Flux1Schnell +from flux_1_schnell.post_processing.image_util import ImageUtil flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell") @@ -18,4 +19,4 @@ image = flux.generate_image( ) ) -image.save("image.png") +ImageUtil.save_image(image, "image.png") diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py index 2a01fd2..51e9599 100644 --- a/src/flux_1_schnell/config/config.py +++ b/src/flux_1_schnell/config/config.py @@ -4,6 +4,7 @@ import logging log = logging.getLogger(__name__) + class Config: precision: mx.Dtype = mx.float16 num_train_steps = 1000 diff --git a/src/flux_1_schnell/post_processing/image_util.py b/src/flux_1_schnell/post_processing/image_util.py index 95f175a..9cc282a 100644 --- a/src/flux_1_schnell/post_processing/image_util.py +++ b/src/flux_1_schnell/post_processing/image_util.py @@ -1,8 +1,13 @@ +import logging +from pathlib import Path + import PIL import mlx.core as mx import numpy as np from PIL import Image +log = logging.getLogger(__name__) + class ImageUtil: @@ -53,3 +58,23 @@ class ImageUtil: def resize(image): image = image.resize((1024, 1024), resample=PIL.Image.LANCZOS) return image + + @staticmethod + def save_image(image: Image.Image, path: str) -> None: + file_path = Path(path) + file_path.parent.mkdir(parents=True, exist_ok=True) + file_name = file_path.stem + file_extension = file_path.suffix + + # If a file already exists, create a new name with a counter + counter = 1 + while file_path.exists(): + new_name = f"{file_name}({counter}){file_extension}" + file_path = file_path.with_name(new_name) + counter += 1 + + try: + image.save(file_path) + log.info(f"Image saved successfully at: {file_path}") + except Exception as e: + log.info(f"Error saving image: {e}") From c50a1fb3d4893eb0b71546bf10bff349413ab7e4 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sun, 18 Aug 2024 21:49:41 +0200 Subject: [PATCH 2/2] Fix height and width argument names --- src/flux_1_schnell/flux.py | 6 +++--- src/flux_1_schnell/models/transformer/transformer.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py index 50dade0..a9feeb0 100644 --- a/src/flux_1_schnell/flux.py +++ b/src/flux_1_schnell/flux.py @@ -61,10 +61,10 @@ class Flux1Schnell: return ImageUtil.to_image(decoded) @staticmethod - def _unpack_latents(latents: mx.array, width: int, height: int) -> mx.array: - latents = mx.reshape(latents, (1, height//16, width//16, 16, 2, 2)) + def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: + latents = mx.reshape(latents, (1, width // 16, height // 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)) + latents = mx.reshape(latents, (1, 16, width // 16 * 2, height // 16 * 2)) return latents def encode(self, path: str) -> mx.array: diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py index 7676885..1d132c9 100644 --- a/src/flux_1_schnell/models/transformer/transformer.py +++ b/src/flux_1_schnell/models/transformer/transformer.py @@ -67,14 +67,14 @@ class Transformer(nn.Module): return noise @staticmethod - def _prepare_latent_image_ids(width: int, height: int) -> mx.array: - latent_height = height // 16 + def _prepare_latent_image_ids(height: int, width: int) -> mx.array: latent_width = width // 16 - latent_image_ids = mx.zeros((latent_height, latent_width, 3)) - latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, latent_height)[:, None]) - latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, latent_width)[None, :]) + latent_height = height // 16 + latent_image_ids = mx.zeros((latent_width, latent_height, 3)) + latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, latent_width)[:, None]) + latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, latent_height)[None, :]) latent_image_ids = mx.repeat(latent_image_ids[None, :], 1, axis=0) - latent_image_ids = mx.reshape(latent_image_ids, (1, latent_height * latent_width, 3)) + latent_image_ids = mx.reshape(latent_image_ids, (1, latent_width * latent_height, 3)) return latent_image_ids @staticmethod