diff --git a/Makefile b/Makefile index b394d48..4e5cf25 100644 --- a/Makefile +++ b/Makefile @@ -93,6 +93,7 @@ check: ensure-ruff .PHONY: test test: ensure-pytest # πŸ—οΈ Running tests... + uv pip install mlx==0.25.0 # Install pinned MLX version specifically for testing $(PYTHON) -m pytest # βœ… Tests completed diff --git a/README.md b/README.md index aff15fb..024a04f 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,3 @@ - ![image](src/mflux/assets/logo.png) *A MLX port of FLUX based on the Huggingface Diffusers implementation.* @@ -35,6 +34,9 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black * [πŸ–ŒοΈ Fill](#%EF%B8%8F-fill) + [Inpainting](#inpainting) + [Outpainting](#outpainting) + * [πŸ” Depth](#-depth) + * [πŸ”„ Redux](#-redux) +- [πŸ•ΉοΈ Controlnet](#%EF%B8%8F-controlnet) - [πŸŽ›οΈ Dreambooth fine-tuning](#-dreambooth-fine-tuning) * [Training configuration](#training-configuration) * [Training example](#training-example) @@ -42,13 +44,12 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black * [Configuration details](#configuration-details) * [Memory issues](#memory-issues) * [Misc](#misc) -- [πŸ•ΉοΈ Controlnet](#%EF%B8%8F-controlnet) - [🚧 Current limitations](#-current-limitations) - [πŸ’‘Workflow tips](#workflow-tips) -- [βœ… TODO](#-todo) - [πŸ”¬ Cool research / features to support](#-cool-research--features-to-support-) - [πŸŒ±β€ Related projects](#-related-projects) -- [License](#license) +- [πŸ™ Acknowledgements](#-acknowledgements) +- [βš–οΈ License](#-license) @@ -869,6 +870,119 @@ mflux-generate-fill \ ⚠️ *Note: Using the Fill tool requires an additional 33.92 GB download from [black-forest-labs/FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev). The download happens automatically on first use.* +#### πŸ” Depth + +Using the depth tool, you can generate high-quality images constrained on a depth map from a reference image. These maps represent the estimated distance of each pixel from the camera, with brighter areas representing objects closer to the camera and darker areas representing objects farther away. + +For state-of-the-art depth extraction, MFLUX now supports the [Depth Pro](https://github.com/apple/ml-depth-pro) model, natively implemented in MLX based on the reference implementation. + +![depth example](src/mflux/assets/depth_example.jpg) +*Original hallway image credit: [Yuliya Matuzava on Unsplash](https://unsplash.com/photos/a-long-hallway-in-a-building-with-arches-and-arches-39NJt9jfGSU)* + +##### Example + +Using the `mflux-generate-depth` we can let MFLUX generate a new image conditioned on the depth map of the input image: + +```bash +mflux-generate-depth \ + --prompt "A hallway made of white carrara marble" \ + -q 8 \ + --height 1680 \ + --width 1200 \ + --image-path "original_image.png" \ + --steps 20 \ + --seed 7102725 +``` + +Instead of automatically generating the depth map based on a reference image, you can instead input the depth map yourself using the `--depth-image-path`, like so: + +```bash +mflux-generate-depth \ + --prompt "A hallway made of white carrara marble" \ + -q 8 \ + --height 1680 \ + --width 1200 \ + --depth-image-path "depth_image.png" \ + --steps 20 +``` + +To generate and export the depth map from an image without triggering the image generation, run the following command: + +```bash +mflux-save-depth --image-path "your_image.jpg" -q 8 +``` + +This will create a depth map and save it with the same name as your image but with a "_depth" suffix (e.g., "your_image_depth.png"). +Quantization is supported for the Depth Pro model, however, output quality can very depend on the input image. + +⚠️ *Note: The Depth Pro model requires an additional 1.9GB download from Apple. The download happens automatically on first use.* + +⚠️ *Note: Using the Depth tool requires an additional 33.92 GB download from [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev). The download happens automatically on first use.* + +#### πŸ”„ Redux + +The Redux tool is an adapter for FLUX.1 models that enables image variation generation similar to the [image-to-image](#-image-to-image) technique. +It can reproduce an input image with slight variations, allowing you to refine and enhance existing images. +Unlike the image-to-image technique which resumes the denoising process mid-way starting from the input image, the redux tool instead embeds the image and joins it with the T5 text encodings. + +![redux example](src/mflux/assets/redux_example.jpg) +*Image credit: [Paris Bilal on Unsplash](https://unsplash.com/photos/a-cat-statue-sitting-on-top-of-a-white-base-Bedd4qHGWCg)* + +##### Example + +The following examples show how to use the Redux tool to generate variations: + +```bash +mflux-generate-redux \ + --prompt "a grey statue of a cat on a white platform in front of a blue background" \ + --redux-image-paths "original.png" \ + --steps 20 \ + --height 1654 \ + --width 1154 \ + -q 8 +``` +There is a tendency for the reference image to dominate over the input prompt. + +⚠️ *Note: Using the Redux tool requires an additional 1.1GB download from [black-forest-labs/FLUX.1-Redux-dev](https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev). The download happens automatically on first use.* + +--- + +### πŸ•ΉοΈ Controlnet + +MFLUX has [Controlnet](https://huggingface.co/docs/diffusers/en/using-diffusers/controlnet) support for an even more fine-grained control +of the image generation. By providing a reference image via `--controlnet-image-path` and a strength parameter via `--controlnet-strength`, you can guide the generation toward the reference image. + +```sh +mflux-generate-controlnet \ + --prompt "A comic strip with a joker in a purple suit" \ + --model dev \ + --steps 20 \ + --seed 1727047657 \ + --height 1066 \ + --width 692 \ + -q 8 \ + --lora-paths "Dark Comic - s0_8 g4.safetensors" \ + --controlnet-image-path "reference.png" \ + --controlnet-strength 0.5 \ + --controlnet-save-canny +``` +![image](src/mflux/assets/controlnet1.jpg) + +*This example combines the controlnet reference image with the LoRA [Dark Comic Flux](https://civitai.com/models/742916/dark-comic-flux)*. + +⚠️ *Note: Controlnet requires an additional one-time download of ~3.58GB of weights from Huggingface. This happens automatically the first time you run the `generate-controlnet` command. +At the moment, the Controlnet used is [InstantX/FLUX.1-dev-Controlnet-Canny](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny), which was trained for the `dev` model. +It can work well with `schnell`, but performance is not guaranteed.* + +⚠️ *Note: The output can be highly sensitive to the controlnet strength and is very much dependent on the reference image. +Too high settings will corrupt the image. A recommended starting point a value like 0.4 and to play around with the strength.* + + +Controlnet can also work well together with [LoRA adapters](#-lora). In the example below the same reference image is used as a controlnet input +with different prompts and LoRA adapters active. + +![image](src/mflux/assets/controlnet2.jpg) + --- ### πŸŽ›οΈ Dreambooth fine-tuning @@ -1027,44 +1141,6 @@ The aim is to also gradually expand the scope of this feature with alternative t - The original fine-tuning script in [Diffusers](https://huggingface.co/docs/diffusers/v0.11.0/en/training/dreambooth) ---- - -### πŸ•ΉοΈ Controlnet - -MFLUX has [Controlnet](https://huggingface.co/docs/diffusers/en/using-diffusers/controlnet) support for an even more fine-grained control -of the image generation. By providing a reference image via `--controlnet-image-path` and a strength parameter via `--controlnet-strength`, you can guide the generation toward the reference image. - -```sh -mflux-generate-controlnet \ - --prompt "A comic strip with a joker in a purple suit" \ - --model dev \ - --steps 20 \ - --seed 1727047657 \ - --height 1066 \ - --width 692 \ - -q 8 \ - --lora-paths "Dark Comic - s0_8 g4.safetensors" \ - --controlnet-image-path "reference.png" \ - --controlnet-strength 0.5 \ - --controlnet-save-canny -``` -![image](src/mflux/assets/controlnet1.jpg) - -*This example combines the controlnet reference image with the LoRA [Dark Comic Flux](https://civitai.com/models/742916/dark-comic-flux)*. - -⚠️ *Note: Controlnet requires an additional one-time download of ~3.58GB of weights from Huggingface. This happens automatically the first time you run the `generate-controlnet` command. -At the moment, the Controlnet used is [InstantX/FLUX.1-dev-Controlnet-Canny](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny), which was trained for the `dev` model. -It can work well with `schnell`, but performance is not guaranteed.* - -⚠️ *Note: The output can be highly sensitive to the controlnet strength and is very much dependent on the reference image. -Too high settings will corrupt the image. A recommended starting point a value like 0.4 and to play around with the strength.* - - -Controlnet can also work well together with [LoRA adapters](#-lora). In the example below the same reference image is used as a controlnet input -with different prompts and LoRA adapters active. - -![image](src/mflux/assets/controlnet2.jpg) - --- ### 🚧 Current limitations @@ -1107,14 +1183,9 @@ See `uv run tools/rename_images.py --help` for full CLI usage help. - When generating multiple images with different seeds, use `--seed` with multiple values or `--auto-seeds` to automatically generate a series of random seeds - Use `--stepwise-image-output-dir` to save intermediate images at each denoising step, which can be useful for debugging or creating animations of the generation process -### βœ… TODO - -- [ ] [FLUX.1 Tools](https://blackforestlabs.ai/flux-1-tools/) - ### πŸ”¬ Cool research / features to support - [ ] [ConceptAttention](https://github.com/helblazer811/ConceptAttention) - [ ] [PuLID](https://github.com/ToTheBeginning/PuLID) -- [ ] [depth based controlnet](https://huggingface.co/InstantX/SD3-Controlnet-Depth) via [ml-depth-pro](https://github.com/apple/ml-depth-pro) or similar? - [ ] [RF-Inversion](https://github.com/filipstrand/mflux/issues/91) - [ ] [catvton-flux](https://github.com/nftblackmagic/catvton-flux) @@ -1125,6 +1196,16 @@ See `uv run tools/rename_images.py --help` for full CLI usage help. - [mflux-fasthtml](https://github.com/anthonywu/mflux-fasthtml) by [@anthonywu](https://github.com/anthonywu) - [mflux-streamlit](https://github.com/elitexp/mflux-streamlit) by [@elitexp](https://github.com/elitexp) -### License +### πŸ™ Acknowledgements + +MFLUX would not be possible without the great work of: + +- The MLX Team for [MLX](https://github.com/ml-explore/mlx) and [MLX examples](https://github.com/ml-explore/mlx-examples) +- Black Forest Labs for the [FLUX project](https://github.com/black-forest-labs/flux) +- Hugging Face for the [Diffusers library implementation of Flux](https://huggingface.co/black-forest-labs/FLUX.1-dev) +- Depth Pro authors for the [Depth Pro model](https://github.com/apple/ml-depth-pro?tab=readme-ov-file#citation) +- The MLX community and all [contributors and testers](https://github.com/filipstrand/mflux/graphs/contributors) + +### βš–οΈ License This project is licensed under the [MIT License](LICENSE). diff --git a/pyproject.toml b/pyproject.toml index 0486fb7..79f4c39 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "tokenizers>=0.20.3; python_version>='3.13'", # transformers -> tokenizers "toml>=0.10.2,<1.0", "torch>=2.3.1,<3.0; python_version<'3.13'", + "torchvision>=0.18.1,<0.22.0", # torch dev builds: pip install --pre --index-url https://download.pytorch.org/whl/nightly "torch>=2.6.0.dev20241106; python_version>='3.13'", "tqdm>=4.66.5,<5.0", @@ -48,6 +49,7 @@ classifiers = [ dev = [ "pytest>=8.3.0,<9.0", "pytest-timer>=1.0,<2.0", + "mlx==0.25.0", # Used ONLY during test runs to ensure deterministic test results ] [project.urls] @@ -58,7 +60,10 @@ mflux-generate = "mflux.generate:main" mflux-generate-controlnet = "mflux.generate_controlnet:main" mflux-generate-in-context = "mflux.generate_in_context:main" mflux-generate-fill = "mflux.generate_fill:main" +mflux-generate-depth = "mflux.generate_depth:main" +mflux-generate-redux = "mflux.generate_redux:main" mflux-save = "mflux.save:main" +mflux-save-depth = "mflux.save_depth:main" mflux-train = "mflux.train:main" [tool.setuptools.packages.find] diff --git a/src/mflux/assets/depth_example.jpg b/src/mflux/assets/depth_example.jpg new file mode 100644 index 0000000..4e319c3 Binary files /dev/null and b/src/mflux/assets/depth_example.jpg differ diff --git a/src/mflux/assets/redux_example.jpg b/src/mflux/assets/redux_example.jpg new file mode 100644 index 0000000..a3a3f09 Binary files /dev/null and b/src/mflux/assets/redux_example.jpg differ diff --git a/src/mflux/callbacks/callback.py b/src/mflux/callbacks/callback.py index efd89e7..4fe5289 100644 --- a/src/mflux/callbacks/callback.py +++ b/src/mflux/callbacks/callback.py @@ -15,6 +15,7 @@ class BeforeLoopCallback(Protocol): latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, + depth_image: PIL.Image.Image | None = None, ) -> None: ... diff --git a/src/mflux/callbacks/callbacks.py b/src/mflux/callbacks/callbacks.py index e415cec..08d07d9 100644 --- a/src/mflux/callbacks/callbacks.py +++ b/src/mflux/callbacks/callbacks.py @@ -14,6 +14,7 @@ class Callbacks: latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, + depth_image: PIL.Image.Image | None = None, ): for subscriber in CallbackRegistry.before_loop_callbacks(): subscriber.call_before_loop( @@ -22,6 +23,7 @@ class Callbacks: latents=latents, config=config, canny_image=canny_image, + depth_image=depth_image, ) @staticmethod diff --git a/src/mflux/callbacks/instances/canny_saver.py b/src/mflux/callbacks/instances/canny_saver.py index 4dcd361..8cb3850 100644 --- a/src/mflux/callbacks/instances/canny_saver.py +++ b/src/mflux/callbacks/instances/canny_saver.py @@ -20,6 +20,7 @@ class CannyImageSaver(BeforeLoopCallback): latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, + depth_image: PIL.Image.Image | None = None, ) -> None: base, ext = os.path.splitext(self.path) ImageUtil.save_image( diff --git a/src/mflux/callbacks/instances/depth_saver.py b/src/mflux/callbacks/instances/depth_saver.py new file mode 100644 index 0000000..0d5a561 --- /dev/null +++ b/src/mflux/callbacks/instances/depth_saver.py @@ -0,0 +1,32 @@ +import os +from pathlib import Path + +import mlx.core as mx +import PIL.Image + +from mflux import ImageUtil +from mflux.callbacks.callback import BeforeLoopCallback +from mflux.config.runtime_config import RuntimeConfig + + +class DepthImageSaver(BeforeLoopCallback): + def __init__(self, path: str): + self.path = Path(path) + + def call_before_loop( + self, + seed: int, + prompt: str, + latents: mx.array, + config: RuntimeConfig, + canny_image: PIL.Image.Image | None = None, + depth_image: PIL.Image.Image | None = None, + ) -> None: + if depth_image is None: + return + + base, ext = os.path.splitext(self.path) + ImageUtil.save_image( + image=depth_image, + path=f"{base}_depth_map{ext}" + ) # fmt: off diff --git a/src/mflux/callbacks/instances/memory_saver.py b/src/mflux/callbacks/instances/memory_saver.py index 932c369..7a3512d 100644 --- a/src/mflux/callbacks/instances/memory_saver.py +++ b/src/mflux/callbacks/instances/memory_saver.py @@ -29,6 +29,7 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback): latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, + depth_image: PIL.Image.Image | None = None, ) -> None: self.peak_memory = mx.metal.get_peak_memory() self._delete_encoders() diff --git a/src/mflux/callbacks/instances/stepwise_handler.py b/src/mflux/callbacks/instances/stepwise_handler.py index 7295ab6..a590baf 100644 --- a/src/mflux/callbacks/instances/stepwise_handler.py +++ b/src/mflux/callbacks/instances/stepwise_handler.py @@ -30,6 +30,7 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, + depth_image: PIL.Image.Image | None = None, ) -> None: self._save_image( step=config.init_time_step, diff --git a/src/mflux/community/in_context_lora/flux_in_context_lora.py b/src/mflux/community/in_context_lora/flux_in_context_lora.py index 90df0bc..f124ef5 100644 --- a/src/mflux/community/in_context_lora/flux_in_context_lora.py +++ b/src/mflux/community/in_context_lora/flux_in_context_lora.py @@ -8,7 +8,7 @@ from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig from mflux.error.exceptions import StopImageGenerationException from mflux.flux.flux_initializer import FluxInitializer -from mflux.latent_creator.latent_creator import Img2Img, LatentCreator +from mflux.latent_creator.latent_creator import LatentCreator from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder from mflux.models.text_encoder.prompt_encoder import PromptEncoder from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder @@ -59,14 +59,10 @@ class Flux1InContextLoRA(nn.Module): # 1. Encode the reference image encoded_image = LatentCreator.encode_image( + vae=self.vae, + image_path=config.image_path, height=config.height, width=config.width, - img2img=Img2Img( - vae=self.vae, - sigmas=config.sigmas, - init_time_step=config.init_time_step, - image_path=config.image_path, - ), ) # 2. Create the initial latents and keep the initial static noise for later blending diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index 99495d2..8fd4817 100644 --- a/src/mflux/config/config.py +++ b/src/mflux/config/config.py @@ -17,6 +17,8 @@ class Config: guidance: float = 4.0, image_path: Path | None = None, image_strength: float | None = None, + depth_image_path: Path | None = None, + redux_image_paths: list[Path] | None = None, masked_image_path: Path | None = None, controlnet_strength: float | None = None, ): @@ -28,5 +30,7 @@ class Config: self.guidance = guidance self.image_path = image_path self.image_strength = image_strength + self.depth_image_path = depth_image_path + self.redux_image_paths = redux_image_paths self.masked_image_path = masked_image_path self.controlnet_strength = controlnet_strength diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py index 7d4b061..1619739 100644 --- a/src/mflux/config/model_config.py +++ b/src/mflux/config/model_config.py @@ -35,11 +35,29 @@ class ModelConfig: def dev_fill() -> "ModelConfig": return AVAILABLE_MODELS["dev-fill"] + @staticmethod + @lru_cache + def dev_depth() -> "ModelConfig": + return AVAILABLE_MODELS["dev-depth"] + + @staticmethod + @lru_cache + def dev_redux() -> "ModelConfig": + return AVAILABLE_MODELS["dev-redux"] + @staticmethod @lru_cache def schnell() -> "ModelConfig": return AVAILABLE_MODELS["schnell"] + def x_embedder_input_dim(self) -> int: + if self.alias == "dev-fill": + return 384 + if self.alias == "dev-depth": + return 128 + else: + return 64 + @staticmethod def from_name( model_name: str, @@ -116,4 +134,24 @@ AVAILABLE_MODELS = { requires_sigma_shift=True, priority=0, ), + "dev-depth": ModelConfig( + alias="dev-depth", + model_name="black-forest-labs/FLUX.1-Depth-dev", + base_model=None, + num_train_steps=1000, + max_sequence_length=512, + supports_guidance=True, + requires_sigma_shift=True, + priority=4, + ), + "dev-redux": ModelConfig( + alias="dev-redux", + model_name="black-forest-labs/FLUX.1-Redux-dev", + base_model=None, + num_train_steps=1000, + max_sequence_length=512, + supports_guidance=True, + requires_sigma_shift=True, + priority=3, + ), } diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index d55073c..85f5eb3 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -55,6 +55,14 @@ class RuntimeConfig: def image_strength(self) -> float | None: return self.config.image_strength + @property + def depth_image_path(self) -> str | None: + return self.config.depth_image_path + + @property + def redux_image_paths(self) -> str | None: + return self.config.redux_image_paths + @property def masked_image_path(self) -> str | None: return self.config.masked_image_path diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index c43a63d..7e8838b 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -61,9 +61,9 @@ class Flux1(nn.Module): width=config.width, img2img=Img2Img( vae=self.vae, + image_path=config.image_path, sigmas=config.sigmas, init_time_step=config.init_time_step, - image_path=config.image_path, ), ) diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index df04f2a..ba918fe 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/src/mflux/flux/flux_initializer.py @@ -1,6 +1,10 @@ from mflux import ModelConfig from mflux.controlnet.transformer_controlnet import TransformerControlnet from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet +from mflux.flux_tools.redux.weight_handler_redux import WeightHandlerRedux +from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer +from mflux.models.redux_encoder.redux_encoder import ReduxEncoder +from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer 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 @@ -26,7 +30,7 @@ class FluxInitializer: lora_names: list[str] | None = None, lora_repo_id: str | None = None, ) -> None: - # 0. Set paths, configs and prompt_cache for later + # 0. Set paths, configs, and prompt_cache for later lora_paths = lora_paths or [] flux_model.prompt_cache = {} flux_model.lora_paths = lora_paths @@ -88,6 +92,65 @@ class FluxInitializer: loras=lora_weights, ) + @staticmethod + def init_depth( + flux_model, + model_config: ModelConfig, + quantize: int | None, + local_path: str | None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + lora_names: list[str] | None = None, + lora_repo_id: str | None = None, + ): + # 1. Start with the same init as regular Flux + FluxInitializer.init( + flux_model=flux_model, + model_config=model_config, + quantize=quantize, + local_path=local_path, + lora_paths=lora_paths, + lora_scales=lora_scales, + lora_names=lora_names, + lora_repo_id=lora_repo_id, + ) + + # 2. Initialize the DepthPro model and assign the weights + flux_model.depth_pro = DepthProInitializer.init() + + @staticmethod + def init_redux( + flux_model, + quantize: int | None, + local_path: str | None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + lora_names: list[str] | None = None, + lora_repo_id: str | None = None, + ): + # 1. Start with the same init as regular Flux dev + FluxInitializer.init( + flux_model=flux_model, + model_config=ModelConfig.dev(), + quantize=quantize, + local_path=local_path, + lora_paths=lora_paths, + lora_scales=lora_scales, + lora_names=lora_names, + lora_repo_id=lora_repo_id, + ) + + # 2. Initialize the redux specific addons + redux_weights = WeightHandlerRedux.load_weights() + flux_model.image_embedder = ReduxEncoder() + flux_model.image_encoder = SiglipVisionTransformer() + WeightUtil.set_redux_weights_and_quantize( + quantize_arg=quantize, + weights=redux_weights, + redux_encoder=flux_model.image_embedder, + siglip_vision_transformer=flux_model.image_encoder, + ) + @staticmethod def init_controlnet( flux_model, @@ -99,7 +162,7 @@ class FluxInitializer: lora_names: list[str] | None = None, lora_repo_id: str | None = None, ) -> None: - # 1. Start with same init as regular Flux + # 1. Start with the same init as regular Flux FluxInitializer.init( flux_model=flux_model, model_config=model_config, diff --git a/src/mflux/flux_tools/depth/__init__.py b/src/mflux/flux_tools/depth/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/flux_tools/depth/depth_util.py b/src/mflux/flux_tools/depth/depth_util.py new file mode 100644 index 0000000..cbac56b --- /dev/null +++ b/src/mflux/flux_tools/depth/depth_util.py @@ -0,0 +1,66 @@ +import logging +import os + +import mlx.core as mx +import PIL.Image + +from mflux.config.runtime_config import RuntimeConfig +from mflux.models.depth_pro.depth_pro import DepthPro +from mflux.models.vae.vae import VAE +from mflux.post_processing.array_util import ArrayUtil +from mflux.post_processing.image_util import ImageUtil + +logger = logging.getLogger(__name__) + + +class DepthUtil: + @staticmethod + def encode_depth_map( + vae: VAE, + depth_pro: DepthPro, + config: RuntimeConfig, + image_path: str | None = None, + depth_image_path: str | None = None, + ) -> (mx.array, PIL.Image.Image): + # 1. Create the depth map or use existing one + depth_image_path, depth_image = DepthUtil.get_or_create_depth_map( + depth_pro=depth_pro, + image_path=image_path, + depth_map_path=depth_image_path, + ) + + # 2. Encode the depth map + scaled_depth_map = ImageUtil.scale_to_dimensions( + image=ImageUtil.load_image(depth_image_path).convert("RGB"), + target_width=config.width, + target_height=config.height, + ) + depth_map_array = ImageUtil.to_array(scaled_depth_map) + encoded_depth = vae.encode(depth_map_array) + depth_latents = ArrayUtil.pack_latents(latents=encoded_depth, height=config.height, width=config.width) + + return depth_latents, depth_image + + @staticmethod + def get_or_create_depth_map( + depth_pro: DepthPro, + image_path: str | None = None, + depth_map_path: str | None = None, + ) -> tuple[str, PIL.Image.Image | None]: + # 1. If a depth map path is provided, use it directly + if depth_map_path: + if not os.path.exists(depth_map_path): + raise FileNotFoundError(f"Depth map file not found: {depth_map_path}") + return depth_map_path, None + if not image_path: + raise ValueError("Either depth_map_path or image_path must be provided") + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + + # 2. Generate a depth map from the image using the DepthPro model + depth_result = depth_pro(image_path) + + # 3. Save the depth map to a file with the same name + _depth suffix + generated_depth_path = str(image_path).rsplit(".", 1)[0] + "_depth.png" + depth_result.depth_image.save(generated_depth_path) + return generated_depth_path, depth_result.depth_image diff --git a/src/mflux/flux_tools/depth/flux_depth.py b/src/mflux/flux_tools/depth/flux_depth.py new file mode 100644 index 0000000..1af8415 --- /dev/null +++ b/src/mflux/flux_tools/depth/flux_depth.py @@ -0,0 +1,157 @@ +import mlx.core as mx +from mlx import nn +from tqdm import tqdm + +from mflux.callbacks.callbacks import Callbacks +from mflux.config.config import Config +from mflux.config.model_config import ModelConfig +from mflux.config.runtime_config import RuntimeConfig +from mflux.error.exceptions import StopImageGenerationException +from mflux.flux.flux_initializer import FluxInitializer +from mflux.flux_tools.depth.depth_util import DepthUtil +from mflux.latent_creator.latent_creator import LatentCreator +from mflux.models.depth_pro.depth_pro import DepthPro +from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder +from mflux.models.text_encoder.prompt_encoder import PromptEncoder +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 + + +class Flux1Depth(nn.Module): + vae: VAE + depth_pro: DepthPro + transformer: Transformer + t5_text_encoder: T5Encoder + clip_text_encoder: CLIPEncoder + + def __init__( + self, + quantize: int | None = None, + local_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + ): + super().__init__() + FluxInitializer.init_depth( + flux_model=self, + model_config=ModelConfig.dev_depth(), + quantize=quantize, + local_path=local_path, + lora_paths=lora_paths, + lora_scales=lora_scales, + ) + + def generate_image( + self, + seed: int, + prompt: str, + config: Config, + ) -> GeneratedImage: + # 0. Create a new runtime config based on the model type and input parameters + config = RuntimeConfig(config, self.model_config) + time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + + # 1. Create the initial latents + latents = LatentCreator.create( + seed=seed, + height=config.height, + width=config.width, + ) + + # 2. Encode the prompt + prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt( + prompt=prompt, + prompt_cache=self.prompt_cache, + t5_tokenizer=self.t5_tokenizer, + clip_tokenizer=self.clip_tokenizer, + t5_text_encoder=self.t5_text_encoder, + clip_text_encoder=self.clip_text_encoder, + ) + + # 3. Encode the depth map + static_depth_map, depth_image = DepthUtil.encode_depth_map( + vae=self.vae, + depth_pro=self.depth_pro, + config=config, + image_path=config.image_path, + depth_image_path=config.depth_image_path, + ) + + # (Optional) Call subscribers for beginning of loop + Callbacks.before_loop( + seed=seed, + prompt=prompt, + latents=latents, + config=config, + depth_image=depth_image, + ) # fmt: off + + for t in time_steps: + try: + # 4.t Concatenate the updated latents with the (static) depth map + hidden_states = mx.concatenate([latents, static_depth_map], axis=-1) + + # 5.t Predict the noise + noise = self.transformer( + t=t, + config=config, + hidden_states=hidden_states, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + ) + + # 6.t Take one denoise step + dt = config.sigmas[t + 1] - config.sigmas[t] + latents += noise * dt + + # (Optional) Call subscribers in-loop + Callbacks.in_loop( + t=t, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=time_steps, + ) # fmt: off + + # (Optional) Evaluate to enable progress tracking + mx.eval(latents) + + except KeyboardInterrupt: # noqa: PERF203 + Callbacks.interruption( + t=t, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=time_steps, + ) + raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}") + + # (Optional) Call subscribers after loop + Callbacks.after_loop( + seed=seed, + prompt=prompt, + latents=latents, + config=config, + ) # fmt: off + + # 7. Decode the latent array and return the image + latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + decoded = self.vae.decode(latents) + return ImageUtil.to_image( + decoded_latents=decoded, + config=config, + seed=seed, + prompt=prompt, + quantization=self.bits, + lora_paths=self.lora_paths, + lora_scales=self.lora_scales, + image_path=config.image_path, + depth_image_path=config.depth_image_path, + generation_time=time_steps.format_dict["elapsed"], + ) diff --git a/src/mflux/flux_tools/fill/flux_fill.py b/src/mflux/flux_tools/fill/flux_fill.py index 503529a..7d56e4d 100644 --- a/src/mflux/flux_tools/fill/flux_fill.py +++ b/src/mflux/flux_tools/fill/flux_fill.py @@ -18,7 +18,6 @@ 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.weights.model_saver import ModelSaver class Flux1Fill(nn.Module): @@ -29,7 +28,6 @@ class Flux1Fill(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, local_path: str | None = None, lora_paths: list[str] | None = None, @@ -38,7 +36,7 @@ class Flux1Fill(nn.Module): super().__init__() FluxInitializer.init( flux_model=self, - model_config=model_config, + model_config=ModelConfig.dev_fill(), quantize=quantize, local_path=local_path, lora_paths=lora_paths, @@ -155,6 +153,3 @@ class Flux1Fill(nn.Module): masked_image_path=config.masked_image_path, generation_time=time_steps.format_dict["elapsed"], ) - - def save_model(self, base_path: str) -> None: - ModelSaver.save_model(self, self.bits, base_path) diff --git a/src/mflux/flux_tools/redux/__init__.py b/src/mflux/flux_tools/redux/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/flux_tools/redux/flux_redux.py b/src/mflux/flux_tools/redux/flux_redux.py new file mode 100644 index 0000000..6756a3b --- /dev/null +++ b/src/mflux/flux_tools/redux/flux_redux.py @@ -0,0 +1,185 @@ +import mlx.core as mx +from mlx import nn +from tqdm import tqdm + +from mflux.callbacks.callbacks import Callbacks +from mflux.config.config import Config +from mflux.config.model_config import ModelConfig +from mflux.config.runtime_config import RuntimeConfig +from mflux.error.exceptions import StopImageGenerationException +from mflux.flux.flux_initializer import FluxInitializer +from mflux.flux_tools.redux.redux_util import ReduxUtil +from mflux.latent_creator.latent_creator import LatentCreator +from mflux.models.redux_encoder.redux_encoder import ReduxEncoder +from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer +from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder +from mflux.models.text_encoder.prompt_encoder import PromptEncoder +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 +from mflux.tokenizer.t5_tokenizer import TokenizerT5 + + +class Flux1Redux(nn.Module): + vae: VAE + image_encoder: SiglipVisionTransformer + image_embedder: ReduxEncoder + transformer: Transformer + t5_text_encoder: T5Encoder + clip_text_encoder: CLIPEncoder + + def __init__( + self, + model_config: ModelConfig, + quantize: int | None = None, + local_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + ): + super().__init__() + FluxInitializer.init_redux( + flux_model=self, + quantize=quantize, + local_path=local_path, + lora_paths=lora_paths, + lora_scales=lora_scales, + ) + + def generate_image( + self, + seed: int, + prompt: str, + config: Config, + ) -> GeneratedImage: + # 0. Create a new runtime config based on the model type and input parameters + config = RuntimeConfig(config, self.model_config) + time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + + # 1. Create the initial latents + latents = LatentCreator.create( + seed=seed, + height=config.height, + width=config.width, + ) + + # 2. Get prompt embeddings by fusing the prompt and image embeddings + prompt_embeds, pooled_prompt_embeds = Flux1Redux._get_prompt_embeddings( + prompt=prompt, + prompt_cache=self.prompt_cache, + t5_tokenizer=self.t5_tokenizer, + clip_tokenizer=self.clip_tokenizer, + t5_text_encoder=self.t5_text_encoder, + clip_text_encoder=self.clip_text_encoder, + image_paths=config.redux_image_paths, + image_encoder=self.image_encoder, + image_embedder=self.image_embedder, + ) # fmt: off + + # (Optional) Call subscribers for beginning of loop + Callbacks.before_loop( + seed=seed, + prompt=prompt, + latents=latents, + config=config, + ) # fmt: off + + for t in time_steps: + try: + # 3.t Predict the noise + noise = self.transformer( + t=t, + config=config, + hidden_states=latents, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + ) + + # 4.t Take one denoise step + dt = config.sigmas[t + 1] - config.sigmas[t] + latents += noise * dt + + # (Optional) Call subscribers in-loop + Callbacks.in_loop( + t=t, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=time_steps, + ) # fmt: off + + # (Optional) Evaluate to enable progress tracking + mx.eval(latents) + + except KeyboardInterrupt: # noqa: PERF203 + Callbacks.interruption( + t=t, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=time_steps, + ) + raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}") + + # (Optional) Call subscribers after loop + Callbacks.after_loop( + seed=seed, + prompt=prompt, + latents=latents, + config=config, + ) # fmt: off + + # 7. Decode the latent array and return the image + latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + decoded = self.vae.decode(latents) + return ImageUtil.to_image( + decoded_latents=decoded, + config=config, + seed=seed, + prompt=prompt, + quantization=self.bits, + lora_paths=self.lora_paths, + lora_scales=self.lora_scales, + redux_image_paths=config.redux_image_paths, + image_strength=config.image_strength, + generation_time=time_steps.format_dict["elapsed"], + ) + + @staticmethod + def _get_prompt_embeddings( + prompt: str, + prompt_cache: dict[str, tuple[mx.array, mx.array]], + t5_tokenizer: TokenizerT5, + clip_tokenizer: TokenizerCLIP, + t5_text_encoder: T5Encoder, + clip_text_encoder: CLIPEncoder, + image_paths: list[str], + image_encoder: SiglipVisionTransformer, + image_embedder: ReduxEncoder, + ) -> (mx.array, mx.array): + # 1. Encode the prompt + prompt_embeds_txt, pooled_prompt_embeds = PromptEncoder.encode_prompt( + prompt=prompt, + prompt_cache=prompt_cache, + t5_tokenizer=t5_tokenizer, + clip_tokenizer=clip_tokenizer, + t5_text_encoder=t5_text_encoder, + clip_text_encoder=clip_text_encoder, + ) + + # 2. Encode the image(s) using the Siglip and Redux encoder + image_embeds = ReduxUtil.embed_images( + image_paths=image_paths, + image_encoder=image_encoder, + image_embedder=image_embedder, + ) # fmt:off + + # 3. Join text embeddings with all image embeddings + prompt_embeds = mx.concatenate([prompt_embeds_txt] + image_embeds, axis=1) + + return prompt_embeds, pooled_prompt_embeds diff --git a/src/mflux/flux_tools/redux/redux_util.py b/src/mflux/flux_tools/redux/redux_util.py new file mode 100644 index 0000000..f5ef347 --- /dev/null +++ b/src/mflux/flux_tools/redux/redux_util.py @@ -0,0 +1,35 @@ +import mlx.core as mx + +from mflux import ImageUtil +from mflux.models.redux_encoder.redux_encoder import ReduxEncoder +from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer + + +class ReduxUtil: + @staticmethod + def embed_images( + image_paths: list[str], + image_encoder: SiglipVisionTransformer, + image_embedder: ReduxEncoder, + ) -> list[mx.array]: # fmt:off + image_embeds_list = [] + for image_path in image_paths: + image_embeds = ReduxUtil._embed_single_image( + image_path=image_path, + image_encoder=image_encoder, + image_embedder=image_embedder, + ) + image_embeds_list.append(image_embeds) + return image_embeds_list + + @staticmethod + def _embed_single_image( + image_path: str, + image_encoder: SiglipVisionTransformer, + image_embedder: ReduxEncoder, + ) -> mx.array: # fmt:off + image = ImageUtil.load_image(image_path).convert("RGB") + image = ImageUtil.preprocess_for_model(image=image) + image_latents, pooler_output = image_encoder(image) + image_embeds = image_embedder(image_latents) + return image_embeds diff --git a/src/mflux/flux_tools/redux/weight_handler_redux.py b/src/mflux/flux_tools/redux/weight_handler_redux.py new file mode 100644 index 0000000..c861cf3 --- /dev/null +++ b/src/mflux/flux_tools/redux/weight_handler_redux.py @@ -0,0 +1,36 @@ +from pathlib import Path + +from huggingface_hub import snapshot_download + +from mflux import ModelConfig +from mflux.weights.weight_handler import MetaData, WeightHandler + + +class WeightHandlerRedux: + def __init__(self, siglip: dict, redux_encoder: dict, meta_data: MetaData): + self.siglip = siglip + self.redux_encoder = redux_encoder + self.meta_data = meta_data + + @staticmethod + def load_weights() -> "WeightHandlerRedux": + root_path = Path(snapshot_download(repo_id=ModelConfig.dev_redux().model_name, allow_patterns=["*.safetensors", "config.json"])) # fmt:off + + siglip, _, _ = WeightHandlerRedux._load_siglip_weights(root_path=root_path) + redux_encoder, _, _ = WeightHandlerRedux._load_redux_encoder_weights(root_path=root_path) + + return WeightHandlerRedux( + siglip=siglip, + redux_encoder=redux_encoder, + meta_data=MetaData(quantization_level=None) + ) # fmt:off + + @staticmethod + def _load_siglip_weights(root_path: Path) -> (dict, int, str | None): + weights, _, _ = WeightHandler.get_weights("image_encoder", root_path) + return weights, _, _ + + @staticmethod + def _load_redux_encoder_weights(root_path: Path) -> (dict, int, str | None): + weights, _, _ = WeightHandler.get_weights("image_embedder", root_path) + return weights, _, _ diff --git a/src/mflux/generate_depth.py b/src/mflux/generate_depth.py new file mode 100644 index 0000000..4860bc0 --- /dev/null +++ b/src/mflux/generate_depth.py @@ -0,0 +1,75 @@ +from mflux import Config, StopImageGenerationException +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.callbacks.instances.depth_saver import DepthImageSaver +from mflux.callbacks.instances.memory_saver import MemorySaver +from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.flux_tools.depth.flux_depth import Flux1Depth +from mflux.ui.cli.parsers import CommandLineParser + + +def main(): + # 0. Parse command line arguments + parser = CommandLineParser(description="Generate an image using the depth tool.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=False) + parser.add_depth_arguments() + parser.add_output_arguments() + args = parser.parse_args() + + # 0. Default to a medium guidance value for depth related tasks. + if args.guidance is None: + args.guidance = 10 + + # 1. Load the model + flux = Flux1Depth( + quantize=args.quantize, + local_path=args.path, + lora_paths=args.lora_paths, + lora_scales=args.lora_scales, + ) + + # 2. Register the optional callbacks + if args.save_depth_map: + CallbackRegistry.register_before_loop(DepthImageSaver(path=args.output)) + if args.stepwise_image_output_dir: + handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) + CallbackRegistry.register_before_loop(handler) + CallbackRegistry.register_in_loop(handler) + CallbackRegistry.register_interrupt(handler) + + memory_saver = None + if args.low_ram: + memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1) + CallbackRegistry.register_before_loop(memory_saver) + CallbackRegistry.register_in_loop(memory_saver) + CallbackRegistry.register_after_loop(memory_saver) + + try: + for seed in args.seed: + # 3. Generate an image for each seed value + image = flux.generate_image( + seed=seed, + prompt=args.prompt, + config=Config( + num_inference_steps=args.steps, + height=args.height, + width=args.width, + guidance=args.guidance, + image_path=args.image_path, + depth_image_path=args.depth_image_path, + ), + ) + + # 4. Save the image + image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) + except StopImageGenerationException as stop_exc: + print(stop_exc) + finally: + if memory_saver: + print(memory_saver.memory_stats()) + + +if __name__ == "__main__": + main() diff --git a/src/mflux/generate_fill.py b/src/mflux/generate_fill.py index fda51d9..ae9e131 100644 --- a/src/mflux/generate_fill.py +++ b/src/mflux/generate_fill.py @@ -1,4 +1,4 @@ -from mflux import Config, ModelConfig, StopImageGenerationException +from mflux import Config, StopImageGenerationException from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.memory_saver import MemorySaver from mflux.callbacks.instances.stepwise_handler import StepwiseHandler @@ -23,7 +23,6 @@ def main(): # 1. Load the model flux = Flux1Fill( - model_config=ModelConfig.dev_fill(), quantize=args.quantize, local_path=args.path, lora_paths=args.lora_paths, @@ -39,7 +38,7 @@ def main(): memory_saver = None if args.low_ram: - memory_saver = MemorySaver(flux, keep_transformer=len(args.seed) > 1) + memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1) CallbackRegistry.register_before_loop(memory_saver) CallbackRegistry.register_in_loop(memory_saver) CallbackRegistry.register_after_loop(memory_saver) diff --git a/src/mflux/generate_in_context.py b/src/mflux/generate_in_context.py index 02482e6..c80810e 100644 --- a/src/mflux/generate_in_context.py +++ b/src/mflux/generate_in_context.py @@ -38,7 +38,7 @@ def main(): memory_saver = None if args.low_ram: - memory_saver = MemorySaver(flux, keep_transformer=len(args.seed) > 1) + memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1) CallbackRegistry.register_before_loop(memory_saver) CallbackRegistry.register_in_loop(memory_saver) CallbackRegistry.register_after_loop(memory_saver) diff --git a/src/mflux/generate_redux.py b/src/mflux/generate_redux.py new file mode 100644 index 0000000..db5bff3 --- /dev/null +++ b/src/mflux/generate_redux.py @@ -0,0 +1,68 @@ +from mflux import Config, ModelConfig, StopImageGenerationException +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.callbacks.instances.memory_saver import MemorySaver +from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.flux_tools.redux.flux_redux import Flux1Redux +from mflux.ui.cli.parsers import CommandLineParser + + +def main(): + # 0. Parse command line arguments + parser = CommandLineParser(description="Generate an image using the fill tool to complete masked areas.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=False) + parser.add_redux_arguments() + parser.add_output_arguments() + args = parser.parse_args() + + # 1. Load the model + flux = Flux1Redux( + model_config=ModelConfig.dev_redux(), + quantize=args.quantize, + local_path=args.path, + lora_paths=args.lora_paths, + lora_scales=args.lora_scales, + ) + + # 2. Register the optional callbacks + if args.stepwise_image_output_dir: + handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) + CallbackRegistry.register_before_loop(handler) + CallbackRegistry.register_in_loop(handler) + CallbackRegistry.register_interrupt(handler) + + memory_saver = None + if args.low_ram: + memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1) + CallbackRegistry.register_before_loop(memory_saver) + CallbackRegistry.register_in_loop(memory_saver) + CallbackRegistry.register_after_loop(memory_saver) + + try: + for seed in args.seed: + # 3. Generate an image for each seed value + image = flux.generate_image( + seed=seed, + prompt=args.prompt, + config=Config( + num_inference_steps=args.steps, + height=args.height, + width=args.width, + guidance=args.guidance, + redux_image_paths=args.redux_image_paths, + ), + ) + + # 4. Save the image + image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) + except StopImageGenerationException as stop_exc: + print(stop_exc) + finally: + if memory_saver: + print(memory_saver.memory_stats()) + + +if __name__ == "__main__": + main() diff --git a/src/mflux/latent_creator/latent_creator.py b/src/mflux/latent_creator/latent_creator.py index 046ee4f..0bd9a8b 100644 --- a/src/mflux/latent_creator/latent_creator.py +++ b/src/mflux/latent_creator/latent_creator.py @@ -57,7 +57,7 @@ class LatentCreator: ) # 2. Encode the image - encoded = LatentCreator.encode_image(height=height, width=width, img2img=img2img) + encoded = LatentCreator.encode_image(vae=img2img.vae, image_path=img2img.image_path, height=height, width=width) # fmt: off latents = ArrayUtil.pack_latents(latents=encoded, height=height, width=width) # 3. Find the appropriate sigma value @@ -71,13 +71,13 @@ class LatentCreator: ) @staticmethod - def encode_image(height: int, width: int, img2img: Img2Img): + def encode_image(vae: VAE, image_path: str, height: int, width: int): scaled_user_image = ImageUtil.scale_to_dimensions( - image=ImageUtil.load_image(img2img.image_path).convert("RGB"), + image=ImageUtil.load_image(image_path).convert("RGB"), target_width=width, target_height=height, ) - encoded = img2img.vae.encode(ImageUtil.to_array(scaled_user_image)) + encoded = vae.encode(ImageUtil.to_array(scaled_user_image)) return encoded @staticmethod diff --git a/src/mflux/models/depth_pro/__init__.py b/src/mflux/models/depth_pro/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/models/depth_pro/conv_utils.py b/src/mflux/models/depth_pro/conv_utils.py new file mode 100644 index 0000000..2f6be4f --- /dev/null +++ b/src/mflux/models/depth_pro/conv_utils.py @@ -0,0 +1,23 @@ +import mlx.core as mx +import mlx.nn as nn + + +class ConvUtils: + @staticmethod + def apply_conv(x: mx.array, conv_module: nn.Module) -> mx.array: + """Apply a convolution with channel format conversion. + + MLX expects channels-last format (B,H,W,C) for convolutions, + but tensors are generally in channels-first format (B,C,H,W). + This helper handles the conversion automatically. + + Args: + x: Input tensor in channels-first format (B,C,H,W) + conv_module: Convolution module to apply + + Returns: + Output tensor in channels-first format (B,C,H,W) + """ + x = mx.transpose(x, (0, 2, 3, 1)) + x = conv_module(x) + return mx.transpose(x, (0, 3, 1, 2)) diff --git a/src/mflux/models/depth_pro/depth_pro.py b/src/mflux/models/depth_pro/depth_pro.py new file mode 100644 index 0000000..2a4ae70 --- /dev/null +++ b/src/mflux/models/depth_pro/depth_pro.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +import torch +from PIL import Image + +from mflux.models.depth_pro.depth_pro_model import DepthProModel +from mflux.post_processing.image_util import ImageUtil + + +@dataclass +class DepthResult: + depth_image: Image.Image + depth_array: mx.array + min_depth: float + max_depth: float + + +class DepthPro(nn.Module): + def __init__(self): + super().__init__() + self.depth_pro_model = DepthProModel() + + def __call__(self, image_path: str | Path, resize: bool = True) -> DepthResult: + input_array, height, width = self._pre_process(image_path) + depth = self.depth_pro_model(input_array) + return self.post_process(depth, height=height, width=width) + + @staticmethod + def _pre_process(image_path): + image = Image.open(image_path).convert("RGB") + input_array = ImageUtil.preprocess_for_depth_pro(image) + input_array = DepthPro._resize(input_array) + return input_array, image.height, image.width + + @staticmethod + def post_process(depth: mx.array, height: int, width: int): + depth_min = mx.min(depth) + depth_max = mx.max(depth) + normalized_depth = (depth - depth_min) / (depth_max - depth_min) + depth_np = np.asarray(normalized_depth.squeeze()) + depth_image = Image.fromarray((depth_np * 255).astype(np.uint8)) + depth_image = depth_image.resize((width, height)) + depth_np = np.array(depth_image) / 255.0 + return DepthResult( + depth_image=depth_image, + depth_array=depth, + min_depth=depth_min.item(), + max_depth=depth_max.item(), + ) + + @staticmethod + def _resize(x: mx.array) -> mx.array: + x_np = np.array(x) + x_torch = torch.from_numpy(x_np) + x_torch = x_torch.unsqueeze(0) + x_torch = torch.nn.functional.interpolate( + x_torch, + size=(1536, 1536), + mode="bilinear", + align_corners=False, + ) + return mx.array(x_torch) diff --git a/src/mflux/models/depth_pro/depth_pro_encoder.py b/src/mflux/models/depth_pro/depth_pro_encoder.py new file mode 100644 index 0000000..e9c093a --- /dev/null +++ b/src/mflux/models/depth_pro/depth_pro_encoder.py @@ -0,0 +1,115 @@ +import math + +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.conv_utils import ConvUtils +from mflux.models.depth_pro.depth_pro_util import DepthProUtil +from mflux.models.depth_pro.dino_v2.dino_vision_transformer import DinoVisionTransformer +from mflux.models.depth_pro.upsample_block import UpSampleBlock + + +class DepthProEncoder(nn.Module): + def __init__(self): + super().__init__() + self.patch_encoder = DinoVisionTransformer() + self.image_encoder = DinoVisionTransformer() + self.upsample_latent0 = UpSampleBlock(dim_in=1024, dim_int=256, dim_out=256, upsample_layers=3) + self.upsample_latent1 = UpSampleBlock(dim_in=1024, dim_out=256, upsample_layers=2) + self.upsample0 = UpSampleBlock(dim_in=1024, dim_out=512, upsample_layers=1) + self.upsample1 = UpSampleBlock(dim_in=1024, dim_out=1024, upsample_layers=1) + self.upsample2 = UpSampleBlock(dim_in=1024, dim_out=1024, upsample_layers=1) + self.upsample_lowres = nn.ConvTranspose2d(in_channels=1024, out_channels=1024, kernel_size=2, stride=2, padding=0, bias=True) # fmt: off + self.fuse_lowres = nn.Conv2d(in_channels=1024 * 2, out_channels=1024, kernel_size=1, stride=1, padding=0, bias=True) # fmt: off + + def __call__(self, x: mx.array) -> list[mx.array]: + # 1. Create the image pyramid + x0, x1, x2 = DepthProUtil.create_pyramid(x) + + # 2: Split to create batched overlapped mini-images at the backbone (BeiT/ViT/Dino) resolution. + x0_patches = DepthProUtil.split(x0, overlap_ratio=0.25) + x1_patches = DepthProUtil.split(x1, overlap_ratio=0.5) + x2_patches = x2 + + # 3: Run the backbone (BeiT) model and get the result of large batch size. + x_pyramid_patches = mx.concatenate((x0_patches, x1_patches, x2_patches), axis=0) + x_pyramid_encodings, backbone_highres_hook0, backbone_highres_hook1 = self.patch_encoder(x_pyramid_patches) + x_pyramid_encodings = DepthProEncoder._reshape_feature(x_pyramid_encodings, width=24, height=24) + + # Calculate indices for splitting + x0_encodings = x_pyramid_encodings[: len(x0_patches)] + x1_encodings = x_pyramid_encodings[len(x0_patches) : len(x0_patches) + len(x1_patches)] + x2_encodings = x_pyramid_encodings[len(x0_patches) + len(x1_patches) :] + + # 4. Merging + x_latent0_encodings = DepthProEncoder._reshape_feature(backbone_highres_hook0, width=24, height=24) + x_latent1_encodings = DepthProEncoder._reshape_feature(backbone_highres_hook1, width=24, height=24) + x_latent0_features = DepthProEncoder._merge(x_latent0_encodings[: 1 * 5 * 5], batch_size=1, padding=3) + x_latent1_features = DepthProEncoder._merge(x_latent1_encodings[: 1 * 5 * 5], batch_size=1, padding=3) + x0_features = DepthProEncoder._merge(x0_encodings, batch_size=1, padding=3) + x1_features = DepthProEncoder._merge(x1_encodings, batch_size=1, padding=6) + x2_features = x2_encodings + + # 5. Upsample feature maps. + x_latent0_features = self.upsample_latent0(x_latent0_features) + x_latent1_features = self.upsample_latent1(x_latent1_features) + x0_features = self.upsample0(x0_features) + x1_features = self.upsample1(x1_features) + x2_features = self.upsample2(x2_features) + + # 6. Apply the image encoder model. + x_global_features, _, _ = self.image_encoder(x2_patches) + x_global_features = DepthProEncoder._reshape_feature(embeddings=x_global_features, width=24, height=24) + x_global_features = ConvUtils.apply_conv(x_global_features, self.upsample_lowres) + x_global_features = mx.concatenate((x2_features, x_global_features), axis=1) + x_global_features = ConvUtils.apply_conv(x_global_features, self.fuse_lowres) + + return [ + x_latent0_features, + x_latent1_features, + x0_features, + x1_features, + x_global_features, + ] + + @staticmethod + def _reshape_feature( + embeddings: mx.array, + width: int, + height: int, + cls_token_offset: int = 1, + ) -> mx.array: + b, hw, c = embeddings.shape + if cls_token_offset > 0: + embeddings = embeddings[:, cls_token_offset:, :] + embeddings = embeddings.reshape(b, height, width, c).transpose(0, 3, 1, 2) + return embeddings + + @staticmethod + def _merge(x: mx.array, batch_size: int, padding: int = 3) -> mx.array: + steps = int(math.sqrt(x.shape[0] // batch_size)) + + idx = 0 + + output_list = [] + for j in range(steps): + output_row_list = [] + for i in range(steps): + output = x[batch_size * idx : batch_size * (idx + 1)] + + if j != 0: + output = output[..., padding:, :] + if i != 0: + output = output[..., :, padding:] + if j != steps - 1: + output = output[..., :-padding, :] + if i != steps - 1: + output = output[..., :, :-padding] + + output_row_list.append(output) + idx += 1 + + output_row = mx.concatenate(output_row_list, axis=-1) + output_list.append(output_row) + output = mx.concatenate(output_list, axis=-2) + return output diff --git a/src/mflux/models/depth_pro/depth_pro_initializer.py b/src/mflux/models/depth_pro/depth_pro_initializer.py new file mode 100644 index 0000000..ba7b8c4 --- /dev/null +++ b/src/mflux/models/depth_pro/depth_pro_initializer.py @@ -0,0 +1,30 @@ +import mlx.nn as nn + +from mflux.models.depth_pro.depth_pro import DepthPro +from mflux.models.depth_pro.weight_handler_depth_pro import WeightHandlerDepthPro + + +class DepthProInitializer: + @staticmethod + def init(quantize: int | None = None) -> DepthPro: + # 1. Initialize the model + depth_pro = DepthPro() + + # 2. Load the weights + depth_pro_weights = WeightHandlerDepthPro.load_weights() + WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent0") + WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent1") + WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample0") + WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample1") + WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample2") + WeightHandlerDepthPro.reposition_head_weights(depth_pro_weights) + WeightHandlerDepthPro.reshape_transposed_convolution_weights(depth_pro_weights) + + # 3. Assign the weights to the model + depth_pro.depth_pro_model.update(depth_pro_weights.weights) + + # 4. Optionally quantize the model + if quantize: + nn.quantize(depth_pro.depth_pro_model, bits=quantize) + + return depth_pro diff --git a/src/mflux/models/depth_pro/depth_pro_model.py b/src/mflux/models/depth_pro/depth_pro_model.py new file mode 100644 index 0000000..d463a35 --- /dev/null +++ b/src/mflux/models/depth_pro/depth_pro_model.py @@ -0,0 +1,19 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.depth_pro_encoder import DepthProEncoder +from mflux.models.depth_pro.fov_head import FOVHead +from mflux.models.depth_pro.multires_conv_decoder import MultiresConvDecoder + + +class DepthProModel(nn.Module): + def __init__(self): + super().__init__() + self.encoder = DepthProEncoder() + self.decoder = MultiresConvDecoder() + self.head = FOVHead() + + def __call__(self, x: mx.array) -> (mx.array, mx.array): + encodings = self.encoder(x) + features = self.decoder(encodings) + return self.head(features) diff --git a/src/mflux/models/depth_pro/depth_pro_util.py b/src/mflux/models/depth_pro/depth_pro_util.py new file mode 100644 index 0000000..eb96d99 --- /dev/null +++ b/src/mflux/models/depth_pro/depth_pro_util.py @@ -0,0 +1,39 @@ +import math + +import mlx.core as mx +import numpy as np +import torch +import torch.nn.functional as F + + +class DepthProUtil: + @staticmethod + def create_pyramid(x: mx.array) -> (mx.array, mx.array, mx.array): + x0 = x + x_np = np.array(x) + x_torch = torch.from_numpy(x_np) + x1_torch = F.interpolate(x_torch, size=None, scale_factor=0.5, mode="bilinear", align_corners=False) + x2_torch = F.interpolate(x_torch, size=None, scale_factor=0.25, mode="bilinear", align_corners=False) + x1 = mx.array(x1_torch.numpy()) + x2 = mx.array(x2_torch.numpy()) + return x0, x1, x2 + + @staticmethod + def split(x: mx.array, overlap_ratio: float = 0.25) -> mx.array: + patch_size = 384 + patch_stride = int(patch_size * (1 - overlap_ratio)) + + image_size = x.shape[-1] + steps = int(math.ceil((image_size - patch_size) / patch_stride)) + 1 + + x_patch_list = [] + for j in range(steps): + j0 = j * patch_stride + j1 = j0 + patch_size + + for i in range(steps): + i0 = i * patch_stride + i1 = i0 + patch_size + x_patch_list.append(x[..., j0:j1, i0:i1]) + + return mx.concatenate(x_patch_list, axis=0) diff --git a/src/mflux/models/depth_pro/dino_v2/__init__.py b/src/mflux/models/depth_pro/dino_v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/models/depth_pro/dino_v2/attention.py b/src/mflux/models/depth_pro/dino_v2/attention.py new file mode 100644 index 0000000..a620e7b --- /dev/null +++ b/src/mflux/models/depth_pro/dino_v2/attention.py @@ -0,0 +1,35 @@ +import mlx.core as mx +from mlx import nn +from mlx.core.fast import scaled_dot_product_attention + + +class Attention(nn.Module): + def __init__( + self, + dim: int = 1024, + head_dim: int = 64, + num_heads: int = 16, + ): + super().__init__() + self.head_dim = head_dim + self.num_heads = num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim, bias=True) + + def __call__(self, x: mx.array) -> mx.array: + # Produce queries, keys and values + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) + qkv = mx.transpose(qkv, (2, 0, 3, 1, 4)) + q, k, v = qkv[0], qkv[1], qkv[2] + + # Scaled dot product attention + scale = 1 / mx.sqrt(q.shape[-1]) + attn_output = scaled_dot_product_attention(q, k, v, scale=scale) + + # Reshape back + attn_output = mx.transpose(attn_output, (0, 2, 1, 3)) + attn_output = mx.reshape(attn_output, (B, -1, self.num_heads * self.head_dim)) + + # Final projection + return self.proj(attn_output) diff --git a/src/mflux/models/depth_pro/dino_v2/dino_vision_transformer.py b/src/mflux/models/depth_pro/dino_v2/dino_vision_transformer.py new file mode 100644 index 0000000..f9172c6 --- /dev/null +++ b/src/mflux/models/depth_pro/dino_v2/dino_vision_transformer.py @@ -0,0 +1,41 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.dino_v2.patch_embed import PatchEmbed +from mflux.models.depth_pro.dino_v2.transformer_block import TransformerBlock + + +class DinoVisionTransformer(nn.Module): + def __init__(self): + super().__init__() + self.cls_token = mx.random.normal(shape=(1, 1, 1024)) + self.pos_embed = mx.random.normal(shape=(1, 577, 1024)) + self.patch_embed = PatchEmbed() + self.blocks = [TransformerBlock() for i in range(24)] + self.norm = nn.LayerNorm(dims=1024, eps=1e-6, bias=True) + + def __call__(self, x: mx.array) -> (mx.array, mx.array, mx.array): + backbone_highres_hook0 = None + backbone_highres_hook1 = None + + x = self.patch_embed(x) + x = self._pos_embed(x) + for i, block in enumerate(self.blocks): + x = block(x) + + # Save intermediary results for later + if i == 5: + backbone_highres_hook0 = x + if i == 11: + backbone_highres_hook1 = x + + x = self.norm(x) + return x, backbone_highres_hook0, backbone_highres_hook1 + + def _pos_embed(self, x: mx.array) -> mx.array: + B, H, W, C = x.shape + x = x.reshape((B, -1, 1024)) + to_cat = [mx.broadcast_to(self.cls_token, (B,) + self.cls_token.shape[1:])] + x = mx.concatenate(to_cat + [x], axis=1) + x = x + self.pos_embed + return x diff --git a/src/mflux/models/depth_pro/dino_v2/layer_scale.py b/src/mflux/models/depth_pro/dino_v2/layer_scale.py new file mode 100644 index 0000000..dd0d036 --- /dev/null +++ b/src/mflux/models/depth_pro/dino_v2/layer_scale.py @@ -0,0 +1,15 @@ +import mlx.core as mx +import mlx.nn as nn + + +class LayerScale(nn.Module): + def __init__( + self, + dims: int, + init_values: float = 1e-5, + ) -> None: + super().__init__() + self.gamma = init_values * mx.ones((dims,)) + + def __call__(self, x: mx.array) -> mx.array: + return x * self.gamma diff --git a/src/mflux/models/depth_pro/dino_v2/mlp.py b/src/mflux/models/depth_pro/dino_v2/mlp.py new file mode 100644 index 0000000..fb68cf3 --- /dev/null +++ b/src/mflux/models/depth_pro/dino_v2/mlp.py @@ -0,0 +1,15 @@ +import mlx.core as mx +import mlx.nn as nn + + +class MLP(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(1024, 4096, bias=True) + self.fc2 = nn.Linear(4096, 1024, bias=True) + + def __call__(self, x: mx.array) -> mx.array: + x = self.fc1(x) + x = nn.gelu(x) + x = self.fc2(x) + return x diff --git a/src/mflux/models/depth_pro/dino_v2/patch_embed.py b/src/mflux/models/depth_pro/dino_v2/patch_embed.py new file mode 100644 index 0000000..01d11f7 --- /dev/null +++ b/src/mflux/models/depth_pro/dino_v2/patch_embed.py @@ -0,0 +1,13 @@ +import mlx.core as mx +import mlx.nn as nn + + +class PatchEmbed(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Conv2d(in_channels=3, out_channels=1024, kernel_size=16, stride=16, bias=True) + + def __call__(self, x: mx.array) -> mx.array: + x = mx.transpose(x, (0, 2, 3, 1)) + x = self.proj(x) + return x diff --git a/src/mflux/models/depth_pro/dino_v2/transformer_block.py b/src/mflux/models/depth_pro/dino_v2/transformer_block.py new file mode 100644 index 0000000..a60039a --- /dev/null +++ b/src/mflux/models/depth_pro/dino_v2/transformer_block.py @@ -0,0 +1,22 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.dino_v2.attention import Attention +from mflux.models.depth_pro.dino_v2.layer_scale import LayerScale +from mflux.models.depth_pro.dino_v2.mlp import MLP + + +class TransformerBlock(nn.Module): + def __init__(self): + super().__init__() + self.norm1 = nn.LayerNorm(dims=1024, eps=1e-6, bias=True) + self.attn = Attention() + self.ls1 = LayerScale(dims=1024, init_values=1e-5) + self.norm2 = nn.LayerNorm(dims=1024, eps=1e-6, bias=True) + self.mlp = MLP() + self.ls2 = LayerScale(dims=1024, init_values=1e-5) + + def __call__(self, x: mx.array) -> mx.array: + x = x + self.ls1(self.attn(self.norm1(x))) + x = x + self.ls2(self.mlp(self.norm2(x))) + return x diff --git a/src/mflux/models/depth_pro/feature_fusion_block_2d.py b/src/mflux/models/depth_pro/feature_fusion_block_2d.py new file mode 100644 index 0000000..ae6b6c5 --- /dev/null +++ b/src/mflux/models/depth_pro/feature_fusion_block_2d.py @@ -0,0 +1,26 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.conv_utils import ConvUtils +from mflux.models.depth_pro.residual_block import ResidualBlock + + +class FeatureFusionBlock2d(nn.Module): + def __init__(self, num_features: int, deconv: bool = False): + super().__init__() + self.use_deconv = deconv + self.resnet1 = ResidualBlock(num_features) + self.resnet2 = ResidualBlock(num_features) + self.deconv = nn.ConvTranspose2d(in_channels=num_features, out_channels=num_features, kernel_size=2, stride=2, padding=0, bias=False) # fmt: off + self.out_conv = nn.Conv2d(in_channels=num_features, out_channels=num_features, kernel_size=1, stride=1, padding=0, bias=True) # fmt: off + + def __call__(self, x0: mx.array, x1: mx.array = None) -> mx.array: + x = x0 + if x1 is not None: + res = self.resnet1(x1) + x = x + res + x = self.resnet2(x) + if self.use_deconv: + x = ConvUtils.apply_conv(x, self.deconv) + x = ConvUtils.apply_conv(x, self.out_conv) + return x diff --git a/src/mflux/models/depth_pro/fov_head.py b/src/mflux/models/depth_pro/fov_head.py new file mode 100644 index 0000000..cb8b6a8 --- /dev/null +++ b/src/mflux/models/depth_pro/fov_head.py @@ -0,0 +1,25 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.conv_utils import ConvUtils + + +class FOVHead(nn.Module): + def __init__(self): + super().__init__() + self.convs = [ + nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, stride=1, padding=1), + nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=2, stride=2, padding=0, bias=True), + nn.Conv2d(in_channels=128, out_channels=32, kernel_size=3, stride=1, padding=1), + nn.Identity(), + nn.Conv2d(in_channels=32, out_channels=1, kernel_size=1, stride=1, padding=0), + ] + + def __call__(self, x: mx.array) -> mx.array: + x = ConvUtils.apply_conv(x, self.convs[0]) + x = ConvUtils.apply_conv(x, self.convs[1]) + x = ConvUtils.apply_conv(x, self.convs[2]) + x = nn.relu(x) + x = ConvUtils.apply_conv(x, self.convs[4]) + x = nn.relu(x) + return x diff --git a/src/mflux/models/depth_pro/multires_conv_decoder.py b/src/mflux/models/depth_pro/multires_conv_decoder.py new file mode 100644 index 0000000..8f34fd4 --- /dev/null +++ b/src/mflux/models/depth_pro/multires_conv_decoder.py @@ -0,0 +1,38 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.conv_utils import ConvUtils +from mflux.models.depth_pro.feature_fusion_block_2d import FeatureFusionBlock2d + + +class MultiresConvDecoder(nn.Module): + def __init__(self): + super().__init__() + self.convs = [ + nn.Identity(), + nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False), + nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False), + nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False), + nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False), + ] + self.fusions = [ + FeatureFusionBlock2d(num_features=256, deconv=False), + FeatureFusionBlock2d(num_features=256, deconv=True), + FeatureFusionBlock2d(num_features=256, deconv=True), + FeatureFusionBlock2d(num_features=256, deconv=True), + FeatureFusionBlock2d(num_features=256, deconv=True), + ] + + def __call__(self, encodings: list[mx.array]) -> tuple[mx.array, mx.array]: + # Process last layer + encodings_last = encodings[4] + features = ConvUtils.apply_conv(encodings_last, self.convs[4]) + features = self.fusions[4](features) + + # Process remaining levels with skip connections + for i in [3, 2, 1, 0]: + enc = encodings[i] + features_i = ConvUtils.apply_conv(enc, self.convs[i]) + features = self.fusions[i](features, features_i) + + return features diff --git a/src/mflux/models/depth_pro/residual_block.py b/src/mflux/models/depth_pro/residual_block.py new file mode 100644 index 0000000..ef1dc7e --- /dev/null +++ b/src/mflux/models/depth_pro/residual_block.py @@ -0,0 +1,36 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.conv_utils import ConvUtils + + +class ResidualBlock(nn.Module): + def __init__(self, num_features: int): + super().__init__() + self.residual = [ + nn.Identity(), + nn.Conv2d( + in_channels=num_features, + out_channels=num_features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ), + nn.Identity(), + nn.Conv2d( + in_channels=num_features, + out_channels=num_features, + kernel_size=3, + stride=1, + padding=1, + bias=True, + ), + ] + + def __call__(self, x: mx.array) -> mx.array: + delta_x = nn.relu(x) + delta_x = ConvUtils.apply_conv(delta_x, self.residual[1]) + delta_x = nn.relu(delta_x) + delta_x = ConvUtils.apply_conv(delta_x, self.residual[3]) + return x + delta_x diff --git a/src/mflux/models/depth_pro/upsample_block.py b/src/mflux/models/depth_pro/upsample_block.py new file mode 100644 index 0000000..3e81c99 --- /dev/null +++ b/src/mflux/models/depth_pro/upsample_block.py @@ -0,0 +1,51 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.depth_pro.conv_utils import ConvUtils + + +class UpSampleBlock(nn.Module): + def __init__( + self, + dim_in: int = 1152, + dim_int: int = 256, + dim_out: int = 256, + upsample_layers: int = 3, + ): + super().__init__() + self.layers = UpSampleBlock._create_layers( + dim_in=dim_in, + dim_int=dim_int, + dim_out=dim_out, + upsample_layers=upsample_layers + ) # fmt: off + + @staticmethod + def _create_layers(dim_in: int, dim_out: int, upsample_layers: int, dim_int: int = None) -> list[nn.Module]: + if dim_int is None: + dim_int = dim_out + + # Create projection layer + layers = [nn.Conv2d(in_channels=dim_in, out_channels=dim_int, kernel_size=1, stride=1, padding=0, bias=False)] + + # Add upsampling layers + layers.extend( + [ + nn.ConvTranspose2d( + in_channels=dim_int if i == 0 else dim_out, + out_channels=dim_out, + kernel_size=2, + stride=2, + padding=0, + bias=False, + ) + for i in range(upsample_layers) + ] + ) + + return layers + + def __call__(self, x: mx.array) -> mx.array: + for layer in self.layers: + x = ConvUtils.apply_conv(x, layer) + return x diff --git a/src/mflux/models/depth_pro/weight_handler_depth_pro.py b/src/mflux/models/depth_pro/weight_handler_depth_pro.py new file mode 100644 index 0000000..2b4a959 --- /dev/null +++ b/src/mflux/models/depth_pro/weight_handler_depth_pro.py @@ -0,0 +1,129 @@ +import logging +import os +import urllib.error +import urllib.request +from pathlib import Path + +import mlx.core as mx +import torch +from mlx.utils import tree_unflatten + +from mflux.weights.weight_handler import MetaData +from mflux.weights.weight_util import WeightUtil + + +class WeightHandlerDepthPro: + def __init__(self, weights: dict, meta_data: MetaData): + self.weights = weights + self.meta_data = meta_data + + @staticmethod + def load_weights() -> "WeightHandlerDepthPro": + model_path = WeightHandlerDepthPro._download_or_get_cached_weights() + pt_weights = torch.load(model_path, map_location="cpu") + weights = WeightHandlerDepthPro._to_mlx_weights(pt_weights) + weights = [WeightUtil.reshape_weights(k, v) for k, v in weights.items()] + weights = WeightUtil.flatten(weights) + weights = tree_unflatten(weights) + return WeightHandlerDepthPro( + weights=weights, + meta_data=MetaData(quantization_level=None) + ) # fmt:off + + @staticmethod + def _to_mlx_weights(pt_weights) -> dict: + mlx_weights = {} + for key, value in pt_weights.items(): + if isinstance(value, torch.Tensor): + mlx_weights[key] = mx.array(value.numpy()) + else: + mlx_weights[key] = value + return mlx_weights + + @staticmethod + def _download_or_get_cached_weights(): + APPLE_MODEL_URL = "https://ml-site.cdn-apple.com/models/depth-pro/depth_pro.pt" + + # 1. Create cache directory for the model + cache_dir = Path(os.path.expanduser("~/.cache/mflux/depth_pro")) + cache_dir.mkdir(parents=True, exist_ok=True) + model_path = cache_dir / "depth_pro.pt" + + # 2. Download if model doesn't exist + if not model_path.exists(): + logging.info("Downloading Depth Pro model from Apple...") + try: + urllib.request.urlretrieve(APPLE_MODEL_URL, model_path) + logging.info(f"Downloaded model to {model_path}") + except (urllib.error.URLError, urllib.error.HTTPError) as e: + logging.error(f"Failed to download model: {e}") + logging.info(f"Please manually download from: {APPLE_MODEL_URL}") + if not model_path.exists(): + raise FileNotFoundError(f"Model file not found at {model_path}") + + return model_path + + @staticmethod + def reposition_encoder_weights(depth_pro_weights, name): + tmp = depth_pro_weights.weights["encoder"][name] + depth_pro_weights.weights["encoder"][name] = {} + depth_pro_weights.weights["encoder"][name]["layers"] = tmp + + @staticmethod + def reposition_head_weights(depth_pro_weights): + tmp = depth_pro_weights.weights["head"] + depth_pro_weights.weights["head"] = {} + depth_pro_weights.weights["head"]["convs"] = tmp + + @staticmethod + def reshape_transposed_convolution_weights(depth_pro_weights): + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 1) + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 2) + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 3) + + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent1", 1) + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent1", 2) + + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample0", 1) + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample1", 1) + WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample2", 1) + + WeightHandlerDepthPro._reshape_upsample_lowres(depth_pro_weights) + + WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 1) + WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 2) + WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 3) + WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 4) + + WeightHandlerDepthPro._reshape_head(depth_pro_weights, 1) + + @staticmethod + def _reshape_upsample(depth_pro_weights, name, layer): + tmp = depth_pro_weights.weights["encoder"][name]["layers"][layer]["weight"] + tmp = WeightHandlerDepthPro._reshape(tmp) + depth_pro_weights.weights["encoder"][name]["layers"][layer]["weight"] = tmp + + @staticmethod + def _reshape_upsample_lowres(depth_pro_weights): + tmp = depth_pro_weights.weights["encoder"]["upsample_lowres"]["weight"] + tmp = WeightHandlerDepthPro._reshape(tmp) + depth_pro_weights.weights["encoder"]["upsample_lowres"]["weight"] = tmp + + @staticmethod + def _reshape_deconv(depth_pro_weights, layer): + tmp = depth_pro_weights.weights["decoder"]["fusions"][layer]["deconv"]["weight"] + tmp = WeightHandlerDepthPro._reshape(tmp) + depth_pro_weights.weights["decoder"]["fusions"][layer]["deconv"]["weight"] = tmp + + @staticmethod + def _reshape_head(depth_pro_weights, layer): + tmp = depth_pro_weights.weights["head"]["convs"][layer]["weight"] + tmp = WeightHandlerDepthPro._reshape(tmp) + depth_pro_weights.weights["head"]["convs"][layer]["weight"] = tmp + + @staticmethod + def _reshape(tensor): + tensor = tensor.transpose(0, 3, 1, 2) + tensor = tensor.transpose(1, 0, 2, 3) + tensor = tensor.transpose(0, 2, 3, 1) + return tensor diff --git a/src/mflux/models/redux_encoder/__init__.py b/src/mflux/models/redux_encoder/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/models/redux_encoder/redux_encoder.py b/src/mflux/models/redux_encoder/redux_encoder.py new file mode 100644 index 0000000..a70ada2 --- /dev/null +++ b/src/mflux/models/redux_encoder/redux_encoder.py @@ -0,0 +1,12 @@ +import mlx.core as mx +from mlx import nn + + +class ReduxEncoder(nn.Module): + def __init__(self): + super().__init__() + self.redux_up = nn.Linear(1152, 4096 * 3) + self.redux_down = nn.Linear(4096 * 3, 4096) + + def __call__(self, x: mx.array) -> mx.array: + return self.redux_down(nn.silu(self.redux_up(x))) diff --git a/src/mflux/models/siglip_vision_transformer/__init__.py b/src/mflux/models/siglip_vision_transformer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/models/siglip_vision_transformer/siglip_encoder.py b/src/mflux/models/siglip_vision_transformer/siglip_encoder.py new file mode 100644 index 0000000..34059ae --- /dev/null +++ b/src/mflux/models/siglip_vision_transformer/siglip_encoder.py @@ -0,0 +1,16 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.siglip_vision_transformer.siglip_encoder_layer import SiglipEncoderLayer + + +class SiglipEncoder(nn.Module): + def __init__(self): + super().__init__() + self.layers = [SiglipEncoderLayer() for _ in range(27)] + + def __call__(self, inputs_embeds: mx.array) -> mx.array: + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states diff --git a/src/mflux/models/siglip_vision_transformer/siglip_encoder_layer.py b/src/mflux/models/siglip_vision_transformer/siglip_encoder_layer.py new file mode 100644 index 0000000..c47a78c --- /dev/null +++ b/src/mflux/models/siglip_vision_transformer/siglip_encoder_layer.py @@ -0,0 +1,29 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.siglip_vision_transformer.siglip_mlp import SiglipMLP +from mflux.models.siglip_vision_transformer.siglip_sdpa_attention import SiglipSdpaAttention + + +class SiglipEncoderLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = SiglipSdpaAttention() + self.layer_norm1 = nn.LayerNorm(1152, eps=1e-6) + self.mlp = SiglipMLP() + self.layer_norm2 = nn.LayerNorm(1152, eps=1e-6) + + def __call__(self, hidden_states: mx.array) -> mx.array: + # Self-attention + residual = hidden_states + hidden_states = self.layer_norm1(hidden_states) + attention_output = self.self_attn(hidden_states) + hidden_states = residual + attention_output + + # MLP + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states diff --git a/src/mflux/models/siglip_vision_transformer/siglip_mlp.py b/src/mflux/models/siglip_vision_transformer/siglip_mlp.py new file mode 100644 index 0000000..2dc7316 --- /dev/null +++ b/src/mflux/models/siglip_vision_transformer/siglip_mlp.py @@ -0,0 +1,15 @@ +import mlx.core as mx +from mlx import nn + + +class SiglipMLP(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(1152, 4304) + self.fc2 = nn.Linear(4304, 1152) + + def __call__(self, hidden_states: mx.array) -> mx.array: + hidden_states = self.fc1(hidden_states) + hidden_states = nn.gelu_approx(hidden_states) # Check that this is equivalent to gelu_pytorch_tanh + hidden_states = self.fc2(hidden_states) + return hidden_states diff --git a/src/mflux/models/siglip_vision_transformer/siglip_multi_head_attention_pooling_head.py b/src/mflux/models/siglip_vision_transformer/siglip_multi_head_attention_pooling_head.py new file mode 100644 index 0000000..dfb1da2 --- /dev/null +++ b/src/mflux/models/siglip_vision_transformer/siglip_multi_head_attention_pooling_head.py @@ -0,0 +1,22 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.siglip_vision_transformer.siglip_mlp import SiglipMLP + + +class SiglipMultiHeadAttentionPoolingHead(nn.Module): + def __init__(self): + super().__init__() + self.probe = mx.random.normal(shape=(1, 1, 1152)) + self.attention = nn.MultiHeadAttention(dims=1152, num_heads=16, bias=True) + self.layernorm = nn.LayerNorm(1152, eps=1e-6) + self.mlp = SiglipMLP() + + def __call__(self, hidden_states: mx.array) -> mx.array: + query = mx.broadcast_to(self.probe, (1, 1, 1152)) + pooled_output = self.attention(query, hidden_states, hidden_states) + residual = pooled_output + pooled_output = self.layernorm(pooled_output) + pooled_output = residual + self.mlp(pooled_output) + pooled_output = pooled_output[:, 0] + return pooled_output diff --git a/src/mflux/models/siglip_vision_transformer/siglip_sdpa_attention.py b/src/mflux/models/siglip_vision_transformer/siglip_sdpa_attention.py new file mode 100644 index 0000000..f2b650b --- /dev/null +++ b/src/mflux/models/siglip_vision_transformer/siglip_sdpa_attention.py @@ -0,0 +1,38 @@ +import mlx.core as mx +from mlx import nn +from mlx.core.fast import scaled_dot_product_attention + + +class SiglipSdpaAttention(nn.Module): + head_dimension = 72 + batch_size = 1 + num_heads = 16 + + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(input_dims=768, output_dims=768) + self.k_proj = nn.Linear(input_dims=768, output_dims=768) + self.v_proj = nn.Linear(input_dims=768, output_dims=768) + self.out_proj = nn.Linear(input_dims=768, output_dims=768) + + def __call__(self, hidden_states: mx.array) -> mx.array: + query = self.q_proj(hidden_states) + key = self.k_proj(hidden_states) + value = self.v_proj(hidden_states) + + query = SiglipSdpaAttention.reshape_and_transpose(query, self.batch_size, self.num_heads, self.head_dimension) + key = SiglipSdpaAttention.reshape_and_transpose(key, self.batch_size, self.num_heads, self.head_dimension) + value = SiglipSdpaAttention.reshape_and_transpose(value, self.batch_size, self.num_heads, self.head_dimension) + + scale = 1 / mx.sqrt(query.shape[-1]) + hidden_states = scaled_dot_product_attention(query, key, value, scale=scale) + + hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3)) + hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension)) + + hidden_states = self.out_proj(hidden_states) + return hidden_states + + @staticmethod + def reshape_and_transpose(x, batch_size, num_heads, head_dim): + return mx.transpose(mx.reshape(x, (batch_size, -1, num_heads, head_dim)), (0, 2, 1, 3)) diff --git a/src/mflux/models/siglip_vision_transformer/siglip_vision_embeddings.py b/src/mflux/models/siglip_vision_transformer/siglip_vision_embeddings.py new file mode 100644 index 0000000..f0a1445 --- /dev/null +++ b/src/mflux/models/siglip_vision_transformer/siglip_vision_embeddings.py @@ -0,0 +1,35 @@ +import mlx.core as mx +from mlx import nn + + +class SiglipVisionEmbeddings(nn.Module): + embed_dim = 1152 + image_size = 384 + patch_size = 14 + + def __init__(self): + super().__init__() + self.patch_embedding = nn.Conv2d( + in_channels=3, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + padding=0, + ) + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + + def __call__(self, pixel_values: mx.array) -> mx.array: + # Get embeddings + pixel_values = mx.transpose(pixel_values, (0, 2, 3, 1)) + embeddings = self.patch_embedding(pixel_values) + embeddings = mx.transpose(embeddings, (0, 3, 1, 2)) + embeddings = embeddings.reshape(1, 1152, 27 * 27) + embeddings = mx.transpose(embeddings, (0, 2, 1)) + + # Get position embeddings + position_ids = mx.arange(self.num_positions) + position_embeddings = self.position_embedding(position_ids) + + return embeddings + position_embeddings diff --git a/src/mflux/models/siglip_vision_transformer/siglip_vision_transformer.py b/src/mflux/models/siglip_vision_transformer/siglip_vision_transformer.py new file mode 100644 index 0000000..a7d9e88 --- /dev/null +++ b/src/mflux/models/siglip_vision_transformer/siglip_vision_transformer.py @@ -0,0 +1,24 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.siglip_vision_transformer.siglip_encoder import SiglipEncoder +from mflux.models.siglip_vision_transformer.siglip_multi_head_attention_pooling_head import ( + SiglipMultiHeadAttentionPoolingHead, +) +from mflux.models.siglip_vision_transformer.siglip_vision_embeddings import SiglipVisionEmbeddings + + +class SiglipVisionTransformer(nn.Module): + def __init__(self): + super().__init__() + self.embeddings = SiglipVisionEmbeddings() + self.encoder = SiglipEncoder() + self.post_layernorm = nn.LayerNorm(dims=1152, eps=1e-6) + self.head = SiglipMultiHeadAttentionPoolingHead() + + def __call__(self, pixel_values: mx.array) -> mx.array: + hidden_states = self.embeddings(pixel_values) + encoder_outputs = self.encoder(hidden_states) + hidden_state = self.post_layernorm(encoder_outputs) + pooler_output = self.head(hidden_state) + return hidden_state, pooler_output diff --git a/src/mflux/models/text_encoder/prompt_encoder.py b/src/mflux/models/text_encoder/prompt_encoder.py index 48fa189..af70160 100644 --- a/src/mflux/models/text_encoder/prompt_encoder.py +++ b/src/mflux/models/text_encoder/prompt_encoder.py @@ -10,7 +10,7 @@ class PromptEncoder: @staticmethod def encode_prompt( prompt: str, - prompt_cache: dict[str, (mx.array, mx.array)], + prompt_cache: dict[str, tuple[mx.array, mx.array]], t5_tokenizer: TokenizerT5, clip_tokenizer: TokenizerCLIP, t5_text_encoder: T5Encoder, diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py index 84e84b0..471c51a 100644 --- a/src/mflux/models/transformer/transformer.py +++ b/src/mflux/models/transformer/transformer.py @@ -27,7 +27,7 @@ class Transformer(nn.Module): ): super().__init__() self.pos_embed = EmbedND() - self.x_embedder = nn.Linear(64, 3072) + self.x_embedder = nn.Linear(model_config.x_embedder_input_dim(), 3072) self.time_text_embed = TimeTextEmbed(model_config=model_config) self.context_embedder = nn.Linear(4096, 3072) self.transformer_blocks = [JointTransformerBlock(i) for i in range(num_transformer_blocks)] diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index eab0d65..4250475 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/src/mflux/post_processing/generated_image.py @@ -28,6 +28,8 @@ class GeneratedImage: image_path: str | pathlib.Path | None = None, image_strength: float | None = None, masked_image_path: str | pathlib.Path | None = None, + depth_image_path: str | pathlib.Path | None = None, + redux_image_paths: list[str] | list[pathlib.Path] | None = None, ): self.image = image self.model_config = model_config @@ -45,6 +47,8 @@ class GeneratedImage: self.image_path = image_path self.image_strength = image_strength self.masked_image_path = masked_image_path + self.depth_image_path = depth_image_path + self.redux_image_paths = redux_image_paths def get_right_half(self) -> "GeneratedImage": # Calculate the coordinates for the right half @@ -69,6 +73,7 @@ class GeneratedImage: image_path=self.image_path, image_strength=self.image_strength, masked_image_path=self.masked_image_path, + depth_image_path=self.depth_image_path, ) def save( @@ -104,6 +109,8 @@ class GeneratedImage: "controlnet_image_path": str(self.controlnet_image_path) if self.controlnet_image_path else None, "controlnet_strength": round(self.controlnet_strength, 2) if self.controlnet_strength else None, "masked_image_path": str(self.masked_image_path) if self.masked_image_path else None, + "depth_image_path": str(self.depth_image_path) if self.depth_image_path else None, + "redux_image_paths": str(self.redux_image_paths) if self.redux_image_paths else None, "prompt": self.prompt, } @@ -113,7 +120,7 @@ class GeneratedImage: if version: return version - # Fallback to installed package version + # Fallback to an installed package version try: return str(importlib.metadata.version("mflux")) except importlib.metadata.PackageNotFoundError: diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index bf65bdf..7729745 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -8,6 +8,7 @@ import numpy as np import piexif import PIL.Image import PIL.ImageDraw +import torchvision.transforms.functional as F from mflux.config.runtime_config import RuntimeConfig from mflux.post_processing.generated_image import GeneratedImage @@ -29,8 +30,10 @@ class ImageUtil: lora_scales: list[float], controlnet_image_path: str | None = None, image_path: str | None = None, + redux_image_paths: list[str] | None = None, image_strength: float | None = None, masked_image_path: str | None = None, + depth_image_path: str | None = None, ) -> GeneratedImage: normalized = ImageUtil._denormalize(decoded_latents) normalized_numpy = ImageUtil._to_numpy(normalized) @@ -52,6 +55,8 @@ class ImageUtil: controlnet_image_path=controlnet_image_path, controlnet_strength=config.controlnet_strength, masked_image_path=masked_image_path, + depth_image_path=depth_image_path, + redux_image_paths=redux_image_paths, ) @staticmethod @@ -260,3 +265,43 @@ class ImageUtil: except Exception as e: # noqa: BLE001 log.error(f"Error embedding metadata: {e}") + + @staticmethod + def preprocess_for_model( + image: PIL.Image.Image, + target_size: tuple = (384, 384), + mean: list = [0.5, 0.5, 0.5], + std: list = [0.5, 0.5, 0.5], + resample: int = PIL.Image.LANCZOS, + ) -> mx.array: + # Resize the image to target size + image = image.resize(target_size, resample=resample) + + # Convert PIL image to numpy array and normalize to [0, 1] + image_np = np.array(image).astype(np.float32) / 255.0 + + # Normalize using specified mean and std + mean_np = np.array(mean) + std_np = np.array(std) + image_np = (image_np - mean_np) / std_np + + # Convert from HWC to CHW format + image_np = image_np.transpose(2, 0, 1) + + # Convert to MLX array and add batch dimension + image_mx = mx.array(image_np) + image_mx = mx.expand_dims(image_mx, axis=0) + + return image_mx + + @staticmethod + def preprocess_for_depth_pro( + image: PIL.Image.Image, + target_size: tuple = (384, 384), + mean: list = [0.5, 0.5, 0.5], + std: list = [0.5, 0.5, 0.5], + resample: int = PIL.Image.LANCZOS, + ) -> mx.array: + image_np = F.to_tensor(image) + tensor = F.normalize(image_np, mean, std, False) + return mx.array(tensor) diff --git a/src/mflux/save.py b/src/mflux/save.py index 0dfe954..cd96b98 100644 --- a/src/mflux/save.py +++ b/src/mflux/save.py @@ -3,24 +3,21 @@ from mflux.ui.cli.parsers import CommandLineParser def main(): + # 0. Parse command line arguments parser = CommandLineParser(description="Save a quantized version of Flux.1 to disk.") # fmt: off parser.add_model_arguments(path_type="save", require_model_arg=True) parser.add_lora_arguments() args = parser.parse_args() - print(f"Saving model {args.model} with quantization level {args.quantize}\n") - + # 1. Load, quantize and save the model flux = Flux1( model_config=ModelConfig.from_name(args.model, base_model=args.base_model), quantize=args.quantize, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) - flux.save_model(args.path) - print(f"Model saved at {args.path}\n") - if __name__ == "__main__": main() diff --git a/src/mflux/save_depth.py b/src/mflux/save_depth.py new file mode 100644 index 0000000..a657133 --- /dev/null +++ b/src/mflux/save_depth.py @@ -0,0 +1,18 @@ +from mflux.flux_tools.depth.depth_util import DepthUtil +from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer +from mflux.ui.cli.parsers import CommandLineParser + + +def main(): + # 0. Parse command line arguments + parser = CommandLineParser(description="Save a depth map of an input image") + parser.add_depth_arguments() + args = parser.parse_args() + + # 1. Create and save the depth map + depth_pro = DepthProInitializer.init(quantize=args.quantize) + DepthUtil.get_or_create_depth_map(depth_pro=depth_pro, image_path=args.image_path) + + +if __name__ == "__main__": + main() diff --git a/src/mflux/train.py b/src/mflux/train.py index 9501373..b09aee0 100644 --- a/src/mflux/train.py +++ b/src/mflux/train.py @@ -5,17 +5,20 @@ from mflux.ui.cli.parsers import CommandLineParser def main(): + # 0. Parse command line arguments parser = CommandLineParser(description="Finetune a LoRA adapter") parser.add_general_arguments() parser.add_model_arguments(require_model_arg=False) parser.add_training_arguments() args = parser.parse_args() + # 1. Initialize the required resources flux, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( config_path=args.train_config, checkpoint_path=args.train_checkpoint, ) + # 2. Start the training process try: DreamBooth.train( flux=flux, diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 9353443..e873732 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -92,6 +92,15 @@ class CommandLineParser(argparse.ArgumentParser): self.add_argument("--image-path", type=Path, required=True, help="Local path to the source image") self.add_argument("--masked-image-path", type=Path, required=True, help="Local path to the mask image") + def add_depth_arguments(self) -> None: + self.add_argument("--image-path", type=Path, required=False, help="Local path to the source image") + self.add_argument("--depth-image-path", type=Path, required=False, help="Local path to the depth image") + self.add_argument("--save-depth-map", action="store_true", required=False, help="If set, save the depth map created from the source image.") + self.add_argument("--quantize", "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)") + + def add_redux_arguments(self) -> None: + self.add_argument("--redux-image-paths", type=Path, nargs="*", required=True, help="Local path to the source image") + def add_output_arguments(self) -> None: self.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.") self.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".") @@ -212,7 +221,7 @@ class CommandLineParser(argparse.ArgumentParser): self.error("--prompt argument required or 'prompt' required in metadata config file") if self.supports_image_generation and namespace.steps is None: - namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, None) + namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, 14) if self.supports_image_outpaint and namespace.image_outpaint_padding is not None: # parse and normalize any acceptable 1,2,3,4-tuple box value to 4-tuple diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py index bfc4a85..2f4a33c 100644 --- a/src/mflux/ui/defaults.py +++ b/src/mflux/ui/defaults.py @@ -6,6 +6,8 @@ MODEL_CHOICES = ["dev", "dev-fill", "schnell"] MODEL_INFERENCE_STEPS = { "dev": 14, "dev-fill": 14, + "dev-depth": 14, + "dev-redux": 14, "schnell": 4, } QUANTIZE_CHOICES = [3, 4, 6, 8] diff --git a/src/mflux/weights/quantization_util.py b/src/mflux/weights/quantization_util.py index 0afa685..a530160 100644 --- a/src/mflux/weights/quantization_util.py +++ b/src/mflux/weights/quantization_util.py @@ -31,9 +31,40 @@ class QuantizationUtil: quantize: int, weights: "WeightHandlerControlnet", transformer_controlnet: nn.Module, - ): + ) -> None: q_level = weights.meta_data.quantization_level if quantize is not None or q_level is not None: bits = int(q_level) if q_level is not None else quantize nn.quantize(transformer_controlnet, bits=bits) + + @staticmethod + def quantize_redux_models( + quantize: int, + weights: "WeightHandler", + redux_encoder: nn.Module, + siglip_vision_transformer: nn.Module, + ) -> None: + q_level = weights.meta_data.quantization_level + + if quantize is not None or q_level is not None: + bits = int(q_level) if q_level is not None else quantize + nn.quantize(redux_encoder, class_predicate=QuantizationUtil.quantization_predicate, bits=bits) + nn.quantize(siglip_vision_transformer, class_predicate=QuantizationUtil.quantization_predicate, bits=bits) + + @staticmethod + def quantization_predicate(path, m): + # 1. Skip Conv2d layers + if isinstance(m, nn.Conv2d): + return False + + # 2. Skip any layer with incompatible dimensions + if hasattr(m, "weight") and hasattr(m.weight, "shape"): + if m.weight.shape == (1152, 4304): + return False + + if m.weight.shape[-1] % 64 != 0: + return False + + # Only quantize layers that have to_quantized method + return hasattr(m, "to_quantized") diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index 830b5e5..7ffaadb 100644 --- a/src/mflux/weights/weight_handler.py +++ b/src/mflux/weights/weight_handler.py @@ -66,12 +66,12 @@ class WeightHandler: @staticmethod def _load_clip_encoder(root_path: Path) -> (dict, int, str | None): - weights, quantization_level, mflux_version = WeightHandler._get_weights("text_encoder", root_path) + weights, quantization_level, mflux_version = WeightHandler.get_weights("text_encoder", root_path) return weights, quantization_level, mflux_version @staticmethod def _load_t5_encoder(root_path: Path) -> (dict, int, str | None): - weights, quantization_level, mflux_version = WeightHandler._get_weights("text_encoder_2", root_path) + weights, quantization_level, mflux_version = WeightHandler.get_weights("text_encoder_2", root_path) # Quantized weights (i.e. ones exported from this project) don't need any post-processing. if quantization_level is not None: @@ -98,7 +98,7 @@ class WeightHandler: @staticmethod def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int, str | None): - weights, quantization_level, mflux_version = WeightHandler._get_weights("transformer", root_path, lora_path) + weights, quantization_level, mflux_version = WeightHandler.get_weights("transformer", root_path, lora_path) if lora_path: if "transformer" not in weights: @@ -126,7 +126,7 @@ class WeightHandler: @staticmethod def _load_vae(root_path: Path) -> (dict, int, str | None): - weights, quantization_level, mflux_version = WeightHandler._get_weights("vae", root_path) + weights, quantization_level, mflux_version = WeightHandler.get_weights("vae", root_path) # Quantized weights (i.e. ones exported from this project) don't need any post-processing. if quantization_level is not None: @@ -142,7 +142,7 @@ class WeightHandler: return weights, quantization_level, mflux_version @staticmethod - def _get_weights( + def get_weights( model_name: str, root_path: Path | None = None, lora_path: str | None = None, diff --git a/src/mflux/weights/weight_util.py b/src/mflux/weights/weight_util.py index 8d603a7..eddf693 100644 --- a/src/mflux/weights/weight_util.py +++ b/src/mflux/weights/weight_util.py @@ -7,6 +7,7 @@ from mflux.weights.quantization_util import QuantizationUtil if TYPE_CHECKING: from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet + from mflux.flux_tools.redux.weight_handler_redux import WeightHandlerRedux from mflux.weights.weight_handler import WeightHandler @@ -83,3 +84,35 @@ class WeightUtil: transformer.update(weights.transformer) t5_text_encoder.update(weights.t5_encoder) clip_text_encoder.update(weights.clip_encoder) + + @staticmethod + def _set_redux_model_weights( + weights: "WeightHandlerRedux", + redux_encoder: nn.Module, + siglip_vision_transformer: nn.Module, + ): + redux_encoder.update(weights.redux_encoder) + siglip_vision_transformer.update(weights.siglip["vision_model"]) + + @staticmethod + def set_redux_weights_and_quantize( + quantize_arg: int | None, + weights: "WeightHandlerRedux", + redux_encoder: nn.Module, + siglip_vision_transformer: nn.Module, + ) -> int | None: + if weights.meta_data.quantization_level is None and quantize_arg is None: + WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer) + return None + + if weights.meta_data.quantization_level is None and quantize_arg is not None: + bits = quantize_arg + WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer) + QuantizationUtil.quantize_redux_models(bits, weights, redux_encoder, siglip_vision_transformer) + return bits + + if weights.meta_data.quantization_level is not None: + bits = weights.meta_data.quantization_level + QuantizationUtil.quantize_redux_models(bits, weights, redux_encoder, siglip_vision_transformer) + WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer) + return bits diff --git a/tests/depth/__init__.py b/tests/depth/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/depth/test_depth_pro.py b/tests/depth/test_depth_pro.py new file mode 100644 index 0000000..c6d8c90 --- /dev/null +++ b/tests/depth/test_depth_pro.py @@ -0,0 +1,38 @@ +import os +from pathlib import Path + +import numpy as np +from PIL import Image + +from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer + + +class TestDepthPro: + def test_depth_pro_generation(self): + # Resolve paths + resource_dir = Path(__file__).parent.parent / "resources" + input_image_path = resource_dir / "reference_controlnet_dev_lora.png" + reference_image_path = resource_dir / "reference_controlnet_dev_lora_depth.png" + output_image_path = resource_dir / "depth_output.png" + + try: + # Initialize DepthPro model + depth_pro = DepthProInitializer.init() + + # Process the image to generate a depth map + depth_result = depth_pro(str(input_image_path)) + + # Save the output image + depth_result.depth_image.save(output_image_path) + + # Assert that the generated depth image matches the reference image + np.testing.assert_array_almost_equal( + np.array(Image.open(output_image_path)), + np.array(Image.open(reference_image_path)), + err_msg=f"Generated depth image doesn't match reference depth image. Check {output_image_path} vs {reference_image_path}", + ) + + finally: + # Cleanup + if os.path.exists(output_image_path): + os.remove(output_image_path) diff --git a/tests/image_generation/helpers/image_generation_depth_test_helper.py b/tests/image_generation/helpers/image_generation_depth_test_helper.py new file mode 100644 index 0000000..68b7ff1 --- /dev/null +++ b/tests/image_generation/helpers/image_generation_depth_test_helper.py @@ -0,0 +1,64 @@ +import os +from pathlib import Path + +import numpy as np +from PIL import Image + +from mflux import Config, ModelConfig +from mflux.flux_tools.depth.flux_depth import Flux1Depth + + +class ImageGeneratorDepthTestHelper: + @staticmethod + def assert_matches_reference_image( + reference_image_path: str, + output_image_path: str, + model_config: ModelConfig, + prompt: str, + steps: int, + seed: int, + height: int = None, + width: int = None, + image_path: str | None = None, + depth_image_path: str | None = None, + ): + # resolve paths + reference_image_path = ImageGeneratorDepthTestHelper.resolve_path(reference_image_path) + output_image_path = ImageGeneratorDepthTestHelper.resolve_path(output_image_path) + depth_image_path = ImageGeneratorDepthTestHelper.resolve_path(depth_image_path) + image_path = ImageGeneratorDepthTestHelper.resolve_path(image_path) + + try: + # given + flux = Flux1Depth(quantize=8) + image = flux.generate_image( + seed=seed, + prompt=prompt, + config=Config( + num_inference_steps=steps, + height=height, + width=width, + image_path=image_path, + depth_image_path=depth_image_path, + ), + ) + + image.save(path=output_image_path, overwrite=True) + + # then + np.testing.assert_array_equal( + np.array(Image.open(output_image_path)), + np.array(Image.open(reference_image_path)), + err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + ) + + finally: + # cleanup + if os.path.exists(output_image_path): + os.remove(output_image_path) + + @staticmethod + def resolve_path(path) -> Path | None: + if path is None: + return None + return Path(__file__).parent.parent.parent / "resources" / path diff --git a/tests/image_generation/helpers/image_generation_fill_test_helper.py b/tests/image_generation/helpers/image_generation_fill_test_helper.py index 7d98f29..ced31f7 100644 --- a/tests/image_generation/helpers/image_generation_fill_test_helper.py +++ b/tests/image_generation/helpers/image_generation_fill_test_helper.py @@ -34,7 +34,6 @@ class ImageGeneratorFillTestHelper: try: # given flux = Flux1Fill( - model_config=model_config, quantize=8, lora_paths=lora_paths, lora_scales=lora_scales, diff --git a/tests/image_generation/helpers/image_generation_redux_test_helper.py b/tests/image_generation/helpers/image_generation_redux_test_helper.py new file mode 100644 index 0000000..3da822d --- /dev/null +++ b/tests/image_generation/helpers/image_generation_redux_test_helper.py @@ -0,0 +1,65 @@ +import os +from pathlib import Path + +import numpy as np +from PIL import Image + +from mflux import Config, ModelConfig +from mflux.flux_tools.redux.flux_redux import Flux1Redux + + +class ImageGeneratorReduxTestHelper: + @staticmethod + def assert_matches_reference_image( + reference_image_path: str, + output_image_path: str, + model_config: ModelConfig, + steps: int, + seed: int, + height: int, + width: int, + prompt: str, + redux_image_path: str, + ): + # resolve paths + reference_image_path = ImageGeneratorReduxTestHelper.resolve_path(reference_image_path) + output_image_path = ImageGeneratorReduxTestHelper.resolve_path(output_image_path) + redux_image_path = ImageGeneratorReduxTestHelper.resolve_path(redux_image_path) + + try: + # given + flux = Flux1Redux( + model_config=model_config, + quantize=8, + ) + + # when + image = flux.generate_image( + seed=seed, + prompt=prompt, + config=Config( + num_inference_steps=steps, + height=height, + width=width, + redux_image_paths=[redux_image_path], + ), + ) + image.save(path=output_image_path) + + # then + np.testing.assert_array_equal( + np.array(Image.open(output_image_path)), + np.array(Image.open(reference_image_path)), + err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}", + ) + + finally: + # cleanup + if os.path.exists(output_image_path): + os.remove(output_image_path) + + @staticmethod + def resolve_path(path) -> Path | None: + if path is None: + return None + return Path(__file__).parent.parent.parent / "resources" / path diff --git a/tests/image_generation/test_generate_image.py b/tests/image_generation/test_generate_image.py index 2ba5f94..4662192 100644 --- a/tests/image_generation/test_generate_image.py +++ b/tests/image_generation/test_generate_image.py @@ -3,12 +3,10 @@ from tests.image_generation.helpers.image_generation_test_helper import ImageGen class TestImageGenerator: - OUTPUT_IMAGE_FILENAME = "output.png" - def test_image_generation_schnell(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_schnell.png", - output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME, + output_image_path="output_schnell.png", model_config=ModelConfig.schnell(), steps=2, seed=42, @@ -20,7 +18,7 @@ class TestImageGenerator: def test_image_generation_dev(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev.png", - output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME, + output_image_path="output_dev.png", model_config=ModelConfig.dev(), steps=15, seed=42, @@ -32,7 +30,7 @@ class TestImageGenerator: def test_image_generation_dev_lora(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_lora.png", - output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME, + output_image_path="output_dev_lora.png", model_config=ModelConfig.dev(), steps=15, seed=42, @@ -46,7 +44,7 @@ class TestImageGenerator: def test_image_generation_dev_multiple_loras(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_lora_multiple.png", - output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME, + output_image_path="output_dev_lora_multiple.png", model_config=ModelConfig.dev(), steps=15, seed=42, @@ -60,13 +58,13 @@ class TestImageGenerator: def test_image_generation_dev_image_to_image(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_image_to_image_result.png", - image_path="reference_dev_lora.png", + output_image_path="output_dev_image_to_image_result.png", image_strength=0.4, - output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME, model_config=ModelConfig.dev(), steps=8, seed=44, height=341, width=768, + image_path="reference_dev_lora.png", prompt="Luxury food photograph of a burger", ) diff --git a/tests/image_generation/test_generate_image_controlnet.py b/tests/image_generation/test_generate_image_controlnet.py index af9131f..97e7bf0 100644 --- a/tests/image_generation/test_generate_image_controlnet.py +++ b/tests/image_generation/test_generate_image_controlnet.py @@ -3,13 +3,12 @@ from tests.image_generation.helpers.image_generation_controlnet_test_helper impo class TestImageGeneratorControlnet: - OUTPUT_IMAGE_FILENAME = "output.png" CONTROLNET_REFERENCE_FILENAME = "controlnet_reference.png" def test_image_generation_schnell_controlnet(self): ImageGeneratorControlnetTestHelper.assert_matches_reference_image( reference_image_path="reference_controlnet_schnell.png", - output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME, + output_image_path="output_controlnet_schnell.png", controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME, model_config=ModelConfig.schnell(), steps=2, @@ -21,7 +20,7 @@ class TestImageGeneratorControlnet: def test_image_generation_dev_controlnet(self): ImageGeneratorControlnetTestHelper.assert_matches_reference_image( reference_image_path="reference_controlnet_dev.png", - output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME, + output_image_path="output_controlnet_dev.png", controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME, model_config=ModelConfig.dev(), steps=15, @@ -33,7 +32,7 @@ class TestImageGeneratorControlnet: def test_image_generation_dev_lora_controlnet(self): ImageGeneratorControlnetTestHelper.assert_matches_reference_image( reference_image_path="reference_controlnet_dev_lora.png", - output_image_path=TestImageGeneratorControlnet.OUTPUT_IMAGE_FILENAME, + output_image_path="output_controlnet_dev_lora.png", controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME, model_config=ModelConfig.dev(), steps=15, diff --git a/tests/image_generation/test_generate_image_depth.py b/tests/image_generation/test_generate_image_depth.py new file mode 100644 index 0000000..6e9a612 --- /dev/null +++ b/tests/image_generation/test_generate_image_depth.py @@ -0,0 +1,19 @@ +from mflux import ModelConfig +from tests.image_generation.helpers.image_generation_depth_test_helper import ImageGeneratorDepthTestHelper + + +class TestImageGeneratorDepth: + SOURCE_IMAGE_FILENAME = "reference_controlnet_dev_lora.png" + + def test_image_generation_with_reference_image(self): + ImageGeneratorDepthTestHelper.assert_matches_reference_image( + reference_image_path="reference_depth_dev_from_image.png", + output_image_path="output_depth_dev_from_image.png", + image_path=TestImageGeneratorDepth.SOURCE_IMAGE_FILENAME, + model_config=ModelConfig.dev_depth(), + steps=15, + seed=42, + height=512, + width=320, + prompt="Cartoon picture of einstein with a cane", + ) diff --git a/tests/image_generation/test_generate_image_fill.py b/tests/image_generation/test_generate_image_fill.py index 6641147..cba1984 100644 --- a/tests/image_generation/test_generate_image_fill.py +++ b/tests/image_generation/test_generate_image_fill.py @@ -3,14 +3,13 @@ from tests.image_generation.helpers.image_generation_fill_test_helper import Ima class TestImageGeneratorFill: - OUTPUT_IMAGE_FILENAME = "output.png" SOURCE_IMAGE_FILENAME = "reference_dev_image_to_image_result.png" MASK_IMAGE_FILENAME = "mask.png" - def test_inpaint(self): + def test_fill(self): ImageGeneratorFillTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_image_to_image_result_inpaint.png", - output_image_path=TestImageGeneratorFill.OUTPUT_IMAGE_FILENAME, + output_image_path="output_dev_image_to_image_result_inpaint.png", model_config=ModelConfig.dev_fill(), steps=15, seed=42, diff --git a/tests/image_generation/test_generate_image_in_context.py b/tests/image_generation/test_generate_image_in_context.py index abdd2cf..584b451 100644 --- a/tests/image_generation/test_generate_image_in_context.py +++ b/tests/image_generation/test_generate_image_in_context.py @@ -3,12 +3,10 @@ from tests.image_generation.helpers.image_generation_in_context_test_helper impo class TestImageGeneratorInContext: - OUTPUT_IMAGE_FILENAME = "output.png" - def test_in_context_lora_identity(self): ImageGeneratorInContextTestHelper.assert_matches_reference_image( reference_image_path="ic_lora_reference_in_context_identity.png", - output_image_path=TestImageGeneratorInContext.OUTPUT_IMAGE_FILENAME, + output_image_path="output_ic_lora_reference_in_context_identity.png", model_config=ModelConfig.dev(), steps=25, seed=42, diff --git a/tests/image_generation/test_generate_image_redux.py b/tests/image_generation/test_generate_image_redux.py new file mode 100644 index 0000000..fde4c7f --- /dev/null +++ b/tests/image_generation/test_generate_image_redux.py @@ -0,0 +1,19 @@ +from mflux import ModelConfig +from tests.image_generation.helpers.image_generation_redux_test_helper import ImageGeneratorReduxTestHelper + + +class TestImageGeneratorRedux: + SOURCE_IMAGE_FILENAME = "reference_dev_image_to_image_result.png" + + def test_image_generation_redux(self): + ImageGeneratorReduxTestHelper.assert_matches_reference_image( + reference_image_path="reference_redux_dev.png", + output_image_path="output_redux_dev.png", + model_config=ModelConfig.dev_redux(), + steps=15, + seed=42, + height=341, + width=768, + prompt="A delicious burger", + redux_image_path=TestImageGeneratorRedux.SOURCE_IMAGE_FILENAME, + ) diff --git a/tests/resources/ic_lora_reference_in_context_identity.png b/tests/resources/ic_lora_reference_in_context_identity.png index 6be7ecd..6698813 100644 Binary files a/tests/resources/ic_lora_reference_in_context_identity.png and b/tests/resources/ic_lora_reference_in_context_identity.png differ diff --git a/tests/resources/reference_controlnet_dev.png b/tests/resources/reference_controlnet_dev.png index 1777975..8c51631 100644 Binary files a/tests/resources/reference_controlnet_dev.png and b/tests/resources/reference_controlnet_dev.png differ diff --git a/tests/resources/reference_controlnet_dev_lora.png b/tests/resources/reference_controlnet_dev_lora.png index ab9912f..cdcde61 100644 Binary files a/tests/resources/reference_controlnet_dev_lora.png and b/tests/resources/reference_controlnet_dev_lora.png differ diff --git a/tests/resources/reference_controlnet_dev_lora_depth.png b/tests/resources/reference_controlnet_dev_lora_depth.png new file mode 100644 index 0000000..1682e28 Binary files /dev/null and b/tests/resources/reference_controlnet_dev_lora_depth.png differ diff --git a/tests/resources/reference_controlnet_schnell.png b/tests/resources/reference_controlnet_schnell.png index 6716c44..b087930 100644 Binary files a/tests/resources/reference_controlnet_schnell.png and b/tests/resources/reference_controlnet_schnell.png differ diff --git a/tests/resources/reference_depth_dev_from_image.png b/tests/resources/reference_depth_dev_from_image.png new file mode 100644 index 0000000..3e71373 Binary files /dev/null and b/tests/resources/reference_depth_dev_from_image.png differ diff --git a/tests/resources/reference_dev.png b/tests/resources/reference_dev.png index 886f9cf..72d6454 100644 Binary files a/tests/resources/reference_dev.png and b/tests/resources/reference_dev.png differ diff --git a/tests/resources/reference_dev_image_to_image_result.png b/tests/resources/reference_dev_image_to_image_result.png index 3e3322e..bf712a1 100644 Binary files a/tests/resources/reference_dev_image_to_image_result.png and b/tests/resources/reference_dev_image_to_image_result.png differ diff --git a/tests/resources/reference_dev_image_to_image_result_inpaint.png b/tests/resources/reference_dev_image_to_image_result_inpaint.png index 7d5608c..f860366 100644 Binary files a/tests/resources/reference_dev_image_to_image_result_inpaint.png and b/tests/resources/reference_dev_image_to_image_result_inpaint.png differ diff --git a/tests/resources/reference_dev_lora.png b/tests/resources/reference_dev_lora.png index 6b6abad..deb9903 100644 Binary files a/tests/resources/reference_dev_lora.png and b/tests/resources/reference_dev_lora.png differ diff --git a/tests/resources/reference_dev_lora_multiple.png b/tests/resources/reference_dev_lora_multiple.png index 932e112..48179fd 100644 Binary files a/tests/resources/reference_dev_lora_multiple.png and b/tests/resources/reference_dev_lora_multiple.png differ diff --git a/tests/resources/reference_redux_dev.png b/tests/resources/reference_redux_dev.png new file mode 100644 index 0000000..5c95711 Binary files /dev/null and b/tests/resources/reference_redux_dev.png differ diff --git a/tests/resources/reference_schnell.png b/tests/resources/reference_schnell.png index f8656c5..87ffbb0 100644 Binary files a/tests/resources/reference_schnell.png and b/tests/resources/reference_schnell.png differ