diff --git a/README.md b/README.md index 3037c7a..00613a6 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black * [🔍 Depth](#-depth) * [🔄 Redux](#-redux) - [🕹ī¸ Controlnet](#%EF%B8%8F-controlnet) +- [🔎 Upscale](#-upscale) - [🎛ī¸ Dreambooth fine-tuning](#-dreambooth-fine-tuning) * [Training configuration](#training-configuration) * [Training example](#training-example) @@ -207,6 +208,10 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 - **`--stepwise-image-output-dir`** (optional, `str`, default: `None`): [EXPERIMENTAL] Output directory to write step-wise images and their final composite image to. This feature may change in future versions. When specified, MFLUX will save an image for each denoising step, allowing you to visualize the generation process from noise to final image. +- **`--vae-tiling`** (optional, flag): Enable VAE tiling to reduce memory usage during the decoding phase. This splits the image into smaller chunks for processing, which can prevent out-of-memory errors when generating high-resolution images. Note that this optimization may occasionally produce a subtle seam in the middle of the image, but it's often worth the tradeoff for being able to generate images that would otherwise cause your system to run out of memory. + +- **`--vae-tiling-split`** (optional, `str`, default: `"horizontal"`): When VAE tiling is enabled, this parameter controls the direction to split the latents. Options are `"horizontal"` (splits into top/bottom) or `"vertical"` (splits into left/right). Use this option to control where potential seams might appear in the final image. + #### 📜 In-Context LoRA Command-Line Arguments The `mflux-generate-in-context` command supports most of the same arguments as `mflux-generate`, with these additional parameters: @@ -1082,6 +1087,41 @@ with different prompts and LoRA adapters active. --- +### 🔎 Upscale + +The upscale tool allows you to increase the resolution of an existing image while maintaining or enhancing its quality and details. It uses a specialized ControlNet model that's trained to intelligently upscale images without introducing artifacts or losing image fidelity. + +Under the hood, this is simply the controlnet pipeline with [jasperai/Flux.1-dev-Controlnet-Upscaler](https://huggingface.co/jasperai/Flux.1-dev-Controlnet-Upscaler), with a small modification that we don't process the image with canny edge detection. + +![upscale example](src/mflux/assets/upscale_example.jpg) +*Image credit: [Kevin Mueller on Unsplash](https://unsplash.com/photos/gray-owl-on-black-background-xvwZJNaiRNo)* + +#### How to use + +```sh +mflux-upscale \ + --prompt "A gray owl on black background" \ + --steps 28 \ + --seed 42 \ + --height 1363 \ + --width 908 \ + -q 8 \ + --controlnet-image-path "low_res_image.png" \ + --controlnet-strength 0.6 +``` + +This will upscale your input image to the specified dimensions. The upscaler works best when increasing the resolution by a factor of 2-4x. + +⚠ī¸ *Note: Depending on the capability of your machine, you might run out of memory when trying to export the image. Try the `--vae-tiling` flag to process the image in smaller chunks, which significantly reduces memory usage during the VAE decoding phase at a minimal cost to performance. This optimization may occasionally produce a subtle seam in the middle of the image, but it's often worth the tradeoff for being able to generate higher resolution images. You can also use `--vae-tiling-split vertical` to change the split direction from the default horizontal (top/bottom) to vertical (left/right).* + +#### Tips for Best Results + +- For optimal results, try to maintain the same aspect ratio as the original image +- Prompting matters: Try to acculturate describe the image when upscaling +- The recommended `--controlnet-strength` is in the range between 0.5 to 0.7 + +--- + ### 🎛ī¸ Dreambooth fine-tuning As of release [v.0.5.0](https://github.com/filipstrand/mflux/releases/tag/v.0.5.0), MFLUX has support for fine-tuning your own LoRA adapters using the [Dreambooth](https://dreambooth.github.io) technique. diff --git a/pyproject.toml b/pyproject.toml index 61f4478..b82c226 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ mflux-generate-redux = "mflux.generate_redux:main" mflux-save = "mflux.save:main" mflux-save-depth = "mflux.save_depth:main" mflux-train = "mflux.train:main" +mflux-upscale = "mflux.upscale:main" [tool.setuptools.packages.find] where = ["src"] diff --git a/src/mflux/assets/upscale_example.jpg b/src/mflux/assets/upscale_example.jpg new file mode 100644 index 0000000..eb58ed6 Binary files /dev/null and b/src/mflux/assets/upscale_example.jpg differ diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py index 1619739..df53186 100644 --- a/src/mflux/config/model_config.py +++ b/src/mflux/config/model_config.py @@ -10,6 +10,7 @@ class ModelConfig: alias: str | None, model_name: str, base_model: str | None, + controlnet_model: str | None, num_train_steps: int, max_sequence_length: int, supports_guidance: bool, @@ -19,6 +20,7 @@ class ModelConfig: self.alias = alias self.model_name = model_name self.base_model = base_model + self.controlnet_model = controlnet_model self.num_train_steps = num_train_steps self.max_sequence_length = max_sequence_length self.supports_guidance = supports_guidance @@ -45,6 +47,21 @@ class ModelConfig: def dev_redux() -> "ModelConfig": return AVAILABLE_MODELS["dev-redux"] + @staticmethod + @lru_cache + def dev_controlnet_canny() -> "ModelConfig": + return AVAILABLE_MODELS["dev-controlnet-canny"] + + @staticmethod + @lru_cache + def schnell_controlnet_canny() -> "ModelConfig": + return AVAILABLE_MODELS["schnell-controlnet-canny"] + + @staticmethod + @lru_cache + def dev_controlnet_upscaler() -> "ModelConfig": + return AVAILABLE_MODELS["dev-controlnet-upscaler"] + @staticmethod @lru_cache def schnell() -> "ModelConfig": @@ -58,6 +75,9 @@ class ModelConfig: else: return 64 + def is_canny(self) -> bool: + return self.alias == "dev-controlnet-canny" or self.alias == "schnell-controlnet-canny" + @staticmethod def from_name( model_name: str, @@ -95,6 +115,7 @@ class ModelConfig: alias=default_base.alias, model_name=model_name, base_model=default_base.model_name, + controlnet_model=default_base.controlnet_model, num_train_steps=default_base.num_train_steps, max_sequence_length=default_base.max_sequence_length, supports_guidance=default_base.supports_guidance, @@ -108,6 +129,7 @@ AVAILABLE_MODELS = { alias="schnell", model_name="black-forest-labs/FLUX.1-schnell", base_model=None, + controlnet_model=None, num_train_steps=1000, max_sequence_length=256, supports_guidance=False, @@ -118,6 +140,7 @@ AVAILABLE_MODELS = { alias="dev", model_name="black-forest-labs/FLUX.1-dev", base_model=None, + controlnet_model=None, num_train_steps=1000, max_sequence_length=512, supports_guidance=True, @@ -128,6 +151,7 @@ AVAILABLE_MODELS = { alias="dev-fill", model_name="black-forest-labs/FLUX.1-Fill-dev", base_model=None, + controlnet_model=None, num_train_steps=1000, max_sequence_length=512, supports_guidance=True, @@ -138,6 +162,7 @@ AVAILABLE_MODELS = { alias="dev-depth", model_name="black-forest-labs/FLUX.1-Depth-dev", base_model=None, + controlnet_model=None, num_train_steps=1000, max_sequence_length=512, supports_guidance=True, @@ -148,10 +173,44 @@ AVAILABLE_MODELS = { alias="dev-redux", model_name="black-forest-labs/FLUX.1-Redux-dev", base_model=None, + controlnet_model=None, num_train_steps=1000, max_sequence_length=512, supports_guidance=True, requires_sigma_shift=True, priority=3, ), + "dev-controlnet-canny": ModelConfig( + alias="dev-controlnet-canny", + model_name="black-forest-labs/FLUX.1-dev", + base_model=None, + controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny", + num_train_steps=1000, + max_sequence_length=512, + supports_guidance=True, + requires_sigma_shift=True, + priority=5, + ), + "schnell-controlnet-canny": ModelConfig( + alias="schnell-controlnet-canny", + model_name="black-forest-labs/FLUX.1-schnell", + base_model=None, + controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny", + num_train_steps=1000, + max_sequence_length=256, + supports_guidance=False, + requires_sigma_shift=False, + priority=6, + ), + "dev-controlnet-upscaler": ModelConfig( + alias="dev-controlnet-upscaler", + model_name="black-forest-labs/FLUX.1-dev", + base_model=None, + controlnet_model="jasperai/Flux.1-dev-Controlnet-Upscaler", + num_train_steps=1000, + max_sequence_length=512, + supports_guidance=False, + requires_sigma_shift=False, + priority=7, + ), } diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index eaf4af1..d0d1d8f 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/src/mflux/controlnet/controlnet_util.py @@ -18,15 +18,18 @@ class ControlnetUtil: height: int, width: int, controlnet_image_path: str, + is_canny: bool, ) -> tuple[mx.array, PIL.Image.Image]: from mflux import ImageUtil control_image = ImageUtil.load_image(controlnet_image_path) control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image) - control_image = ControlnetUtil._preprocess_canny(control_image) + if is_canny: + control_image = ControlnetUtil._preprocess_canny(control_image) controlnet_cond = ImageUtil.to_array(control_image) controlnet_cond = vae.encode(controlnet_cond) - controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor + if is_canny: + controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=height, width=width) return controlnet_cond, control_image diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index addd8b9..c3fd8d8 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -65,6 +65,7 @@ class Flux1Controlnet(nn.Module): height=config.height, width=config.width, controlnet_image_path=controlnet_image_path, + is_canny=self.model_config.is_canny(), ) # 2. Create the initial latents diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py index 93988df..411aee7 100644 --- a/src/mflux/controlnet/weight_handler_controlnet.py +++ b/src/mflux/controlnet/weight_handler_controlnet.py @@ -8,8 +8,6 @@ from mlx.utils import tree_unflatten from mflux.weights.weight_handler import MetaData from mflux.weights.weight_util import WeightUtil -CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny" - class WeightHandlerControlnet: def __init__(self, meta_data: MetaData, config: dict, controlnet_transformer: dict | None = None): @@ -18,8 +16,8 @@ class WeightHandlerControlnet: self.config = config @staticmethod - def load_controlnet_transformer() -> "WeightHandlerControlnet": - controlnet_path = Path(snapshot_download(repo_id=CONTROLNET_ID, allow_patterns=["*.safetensors", "config.json"])) # fmt:off + def load_controlnet_transformer(controlnet_model: str) -> "WeightHandlerControlnet": + controlnet_path = Path(snapshot_download(repo_id=controlnet_model, allow_patterns=["*.safetensors", "config.json"])) # fmt:off file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors")) quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level") weights = list(mx.load(str(file)).items()) diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index 89d9a64..6a71362 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/src/mflux/flux/flux_initializer.py @@ -175,7 +175,9 @@ class FluxInitializer: ) # 2. Apply ControlNet-specific initialization - weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer() + weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer( + controlnet_model=model_config.controlnet_model + ) flux_model.transformer_controlnet = TransformerControlnet( model_config=model_config, num_transformer_blocks=weights_controlnet.num_transformer_blocks(), diff --git a/src/mflux/generate.py b/src/mflux/generate.py index df7cf1f..432e4ff 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -62,6 +62,11 @@ def _register_callbacks(args: Namespace, flux: Flux1) -> MemorySaver | None: battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) CallbackRegistry.register_before_loop(battery_saver) + # VAE Tiling + if args.vae_tiling: + flux.vae.decoder.enable_tiling = True + flux.vae.decoder.split_direction = args.vae_tiling_split + # Stepwise Handler if args.stepwise_image_output_dir: handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 3d5c1d6..badec90 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -24,7 +24,7 @@ def main(): # 1. Load the model flux = Flux1Controlnet( - model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model), + model_config=_get_controlnet_model_config(args.model), quantize=args.quantize, local_path=args.path, lora_paths=args.lora_paths, @@ -59,11 +59,22 @@ def main(): print(memory_saver.memory_stats()) +def _get_controlnet_model_config(model_name: str) -> ModelConfig: + if model_name == "schnell": + return ModelConfig.schnell_controlnet_canny() + return ModelConfig.dev_controlnet_canny() + + def _register_callbacks(args: Namespace, flux: Flux1Controlnet) -> MemorySaver | None: # Battery saver battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) CallbackRegistry.register_before_loop(battery_saver) + # VAE Tiling + if args.vae_tiling: + flux.vae.decoder.enable_tiling = True + flux.vae.decoder.split_direction = args.vae_tiling_split + # Canny Image Saver if args.controlnet_save_canny: canny_image_saver = CannyImageSaver(path=args.output) diff --git a/src/mflux/generate_depth.py b/src/mflux/generate_depth.py index b4a6339..4447fbb 100644 --- a/src/mflux/generate_depth.py +++ b/src/mflux/generate_depth.py @@ -68,6 +68,11 @@ def _register_callbacks(args: Namespace, flux: Flux1Depth) -> MemorySaver | None battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) CallbackRegistry.register_before_loop(battery_saver) + # VAE Tiling + if args.vae_tiling: + flux.vae.decoder.enable_tiling = True + flux.vae.decoder.split_direction = args.vae_tiling_split + # Depth Image Saver if args.save_depth_map: depth_image_saver = DepthImageSaver(path=args.output) diff --git a/src/mflux/generate_fill.py b/src/mflux/generate_fill.py index 7e9466c..84dcbc8 100644 --- a/src/mflux/generate_fill.py +++ b/src/mflux/generate_fill.py @@ -67,6 +67,11 @@ def _register_callbacks(args: Namespace, flux: Flux1Fill) -> MemorySaver | None: battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) CallbackRegistry.register_before_loop(battery_saver) + # VAE Tiling + if args.vae_tiling: + flux.vae.decoder.enable_tiling = True + flux.vae.decoder.split_direction = args.vae_tiling_split + # Stepwise Handler if args.stepwise_image_output_dir: handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) diff --git a/src/mflux/generate_in_context.py b/src/mflux/generate_in_context.py index 6472590..0c6584e 100644 --- a/src/mflux/generate_in_context.py +++ b/src/mflux/generate_in_context.py @@ -75,6 +75,11 @@ def _register_callbacks(args: Namespace, flux: Flux1InContextLoRA) -> MemorySave battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) CallbackRegistry.register_before_loop(battery_saver) + # VAE Tiling + if args.vae_tiling: + flux.vae.decoder.enable_tiling = True + flux.vae.decoder.split_direction = "vertical" + # Stepwise Handler if args.stepwise_image_output_dir: handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) diff --git a/src/mflux/generate_redux.py b/src/mflux/generate_redux.py index 8265b46..e2340ba 100644 --- a/src/mflux/generate_redux.py +++ b/src/mflux/generate_redux.py @@ -71,6 +71,11 @@ def _register_callbacks(args: Namespace, flux: Flux1Redux) -> MemorySaver | None battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) CallbackRegistry.register_before_loop(battery_saver) + # VAE Tiling + if args.vae_tiling: + flux.vae.decoder.enable_tiling = True + flux.vae.decoder.split_direction = args.vae_tiling_split + # Stepwise Handler if args.stepwise_image_output_dir: handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir) diff --git a/src/mflux/models/vae/decoder/decoder.py b/src/mflux/models/vae/decoder/decoder.py index 1196dd0..87e8bef 100644 --- a/src/mflux/models/vae/decoder/decoder.py +++ b/src/mflux/models/vae/decoder/decoder.py @@ -11,7 +11,7 @@ from mflux.models.vae.decoder.up_block_4 import UpBlock4 class Decoder(nn.Module): - def __init__(self): + def __init__(self, enable_tiling: bool = False, split_direction: str = "horizontal"): super().__init__() self.conv_in = ConvIn() self.mid_block = UnetMidBlock() @@ -23,13 +23,78 @@ class Decoder(nn.Module): ] self.conv_norm_out = ConvNormOut() self.conv_out = ConvOut() + self.enable_tiling = enable_tiling + self.split_direction = split_direction def __call__(self, latents: mx.array) -> mx.array: latents = self.conv_in(latents) latents = self.mid_block(latents) - for up_block in self.up_blocks: - latents = up_block(latents) + for i, up_block in enumerate(self.up_blocks): + latents = self._apply_up_block(i, up_block, latents) latents = self.conv_norm_out(latents) latents = nn.silu(latents) latents = self.conv_out(latents) return latents + + def _apply_up_block( + self, + block_id: int, + up_block: nn.Module, + latents: mx.array, + ) -> mx.array: + if self.enable_tiling: + return Decoder._apply_up_block_with_tiling(block_id, up_block, latents, self.split_direction) + else: + return up_block(latents) + + @staticmethod + def _apply_up_block_with_tiling( + block_id: int, + up_block: nn.Module, + latents: mx.array, + split_direction: str = "horizontal", + ) -> mx.array: + # Turns out that the third block is the first to cause OOM issues for larger resolutions. + # This trick is a "lossy optimization" that splits the latents and processes them separately. + # The resulting image may have a slightly visible seam at the split point. + if block_id == 2: + latents = Decoder._process_block_3_in_tiles(latents, up_block, split_direction) + else: + latents = up_block(latents) + return latents + + @staticmethod + def _process_block_3_in_tiles( + latents: mx.array, up_block: nn.Module, split_direction: str = "horizontal" + ) -> mx.array: + B, C, H, W = latents.shape + if split_direction == "horizontal": + return Decoder._process_horizontal(H, latents, up_block) + else: + return Decoder._process_vertical(W, latents, up_block) + + @staticmethod + def _process_vertical(width: int, latents: mx.array, up_block: nn.Module) -> mx.array: + # 1. Tile left and right + left_tile_input = latents[:, :, :, : width // 2] + right_tile_input = latents[:, :, :, width // 2 :] + + # 2. Process each tile individually + processed_left_tile = up_block(left_tile_input) + processed_right_tile = up_block(right_tile_input) + + # 3. Concatenate along width (axis=3) + return mx.concatenate([processed_left_tile, processed_right_tile], axis=3) + + @staticmethod + def _process_horizontal(height: int, latents: mx.array, up_block: nn.Module) -> mx.array: + # 1. Tile top and bottom + top_tile_input = latents[:, :, : height // 2, :] + bottom_tile_input = latents[:, :, height // 2 :, :] + + # 2. Process each tile individually + processed_top_tile = up_block(top_tile_input) + processed_bottom_tile = up_block(bottom_tile_input) + + # 3. Concatenate along height (axis=2) + return mx.concatenate([processed_top_tile, processed_bottom_tile], axis=2) diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 0501c7e..55a27b4 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -115,7 +115,7 @@ class ImageUtil: @staticmethod def load_image(path: str | Path) -> PIL.Image.Image: - return PIL.Image.open(path) + return PIL.Image.open(path).convert("RGB") @staticmethod def expand_image( diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 30c5c60..e12f88d 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -44,6 +44,8 @@ class CommandLineParser(argparse.ArgumentParser): def add_general_arguments(self) -> None: self.add_argument("--battery-percentage-stop-limit", "-B", type=lambda v: max(min(int(v), 99), 1), default=ui_defaults.BATTERY_PERCENTAGE_STOP_LIMIT, help=f"On Macs powered by battery, stop image generation when battery reaches this percentage. Default: {ui_defaults.BATTERY_PERCENTAGE_STOP_LIMIT}") self.add_argument("--low-ram", action="store_true", help="Enable low-RAM mode to reduce memory usage (may impact performance).") + self.add_argument("--vae-tiling", action="store_true", help="Enable VAE tiling to reduce memory usage at the cost of a potential seam") + self.add_argument("--vae-tiling-split", type=str, choices=["horizontal", "vertical"], default="horizontal", help="Direction to split the latents when using VAE tiling (horizontal = top/bottom, vertical = left/right). Default is horizontal.") def add_model_arguments(self, path_type: t.Literal["load", "save"] = "load", require_model_arg: bool = True) -> None: self.require_model_arg = require_model_arg diff --git a/src/mflux/upscale.py b/src/mflux/upscale.py new file mode 100644 index 0000000..1fc7c8b --- /dev/null +++ b/src/mflux/upscale.py @@ -0,0 +1,88 @@ +from argparse import Namespace + +from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationException +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.callbacks.instances.battery_saver import BatterySaver +from mflux.callbacks.instances.memory_saver import MemorySaver +from mflux.callbacks.instances.stepwise_handler import StepwiseHandler +from mflux.error.exceptions import PromptFileReadError +from mflux.ui.cli.parsers import CommandLineParser +from mflux.ui.prompt_utils import get_effective_prompt + + +def main(): + # 0. Parse command line arguments + parser = CommandLineParser(description="Upscale an image.") + 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_controlnet_arguments() + parser.add_output_arguments() + args = parser.parse_args() + + # 1. Load the model + flux = Flux1Controlnet( + model_config=ModelConfig.dev_controlnet_upscaler(), + quantize=args.quantize, + local_path=args.path, + lora_paths=args.lora_paths, + lora_scales=args.lora_scales, + ) + + # 2. Register the optional callbacks + memory_saver = _register_callbacks(args=args, flux=flux) + + try: + for seed in args.seed: + # 3. Generate an upscaled image for each seed value + image = flux.generate_image( + seed=seed, + prompt=get_effective_prompt(args), + controlnet_image_path=args.controlnet_image_path, + config=Config( + num_inference_steps=args.steps, + height=args.height, + width=args.width, + controlnet_strength=args.controlnet_strength, + ), + ) + + # 4. Save the image + image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) + finally: + if memory_saver: + print(memory_saver.memory_stats()) + + +def _register_callbacks(args: Namespace, flux: Flux1Controlnet) -> MemorySaver | None: + # Battery saver + battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) + CallbackRegistry.register_before_loop(battery_saver) + + # VAE Tiling + if args.vae_tiling: + flux.vae.decoder.enable_tiling = True + flux.vae.decoder.split_direction = args.vae_tiling_split + + # Stepwise Handler + 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 + 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) + return memory_saver + + +if __name__ == "__main__": + main() diff --git a/tests/callbacks/test_battery_saver.py b/tests/callbacks/test_battery_saver.py index 94acb4f..6531508 100644 --- a/tests/callbacks/test_battery_saver.py +++ b/tests/callbacks/test_battery_saver.py @@ -6,24 +6,6 @@ from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_pe from mflux.error.exceptions import StopImageGenerationException -def test_get_battery_percentage_success(): - """Test that the battery percentage is correctly extracted when subprocess returns a valid output.""" - with patch("subprocess.run") as mock_run: - # Set up mock to return a sample battery status output - mock_result = MagicMock() - mock_result.stdout = "Now drawing from 'Battery Power'\n -InternalBattery-0 (id=1234567) 42%;\n" - mock_run.return_value = mock_result - - # Call the function - percentage = get_battery_percentage() - - # Verify subprocess.run was called with the correct arguments - mock_run.assert_called_once_with(["pmset", "-g", "batt"], capture_output=True, text=True, check=True) - - # Assert the function properly extracted the battery percentage - assert percentage == 42 - - def test_get_battery_percentage_while_charging(): """Test that the function returns None when the output doesn't match the expected pattern.""" with patch("subprocess.run") as mock_run: diff --git a/tests/image_generation/helpers/__init__.py b/tests/image_generation/helpers/__init__.py index d403634..e69de29 100644 --- a/tests/image_generation/helpers/__init__.py +++ b/tests/image_generation/helpers/__init__.py @@ -1,9 +0,0 @@ -from tests.image_generation.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper -from tests.image_generation.helpers.image_generation_in_context_test_helper import ImageGeneratorInContextTestHelper -from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper - -__all__ = [ - "ImageGeneratorTestHelper", - "ImageGeneratorControlnetTestHelper", - "ImageGeneratorInContextTestHelper", -] diff --git a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py index d132fd1..6a8f0b0 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py @@ -17,6 +17,8 @@ class ImageGeneratorControlnetTestHelper: prompt: str, steps: int, seed: int, + height: int, + width: int, controlnet_strength: float, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, @@ -43,8 +45,8 @@ class ImageGeneratorControlnetTestHelper: controlnet_image_path=controlnet_image_path, config=Config( num_inference_steps=steps, - height=768, - width=493, + height=height, + width=width, controlnet_strength=controlnet_strength, ), ) diff --git a/tests/image_generation/test_generate_image_controlnet.py b/tests/image_generation/test_generate_image_controlnet.py index 97e7bf0..0eb7d50 100644 --- a/tests/image_generation/test_generate_image_controlnet.py +++ b/tests/image_generation/test_generate_image_controlnet.py @@ -10,9 +10,11 @@ class TestImageGeneratorControlnet: reference_image_path="reference_controlnet_schnell.png", output_image_path="output_controlnet_schnell.png", controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME, - model_config=ModelConfig.schnell(), + model_config=ModelConfig.schnell_controlnet_canny(), steps=2, seed=43, + height=768, + width=493, prompt="The joker with a hat and a cane", controlnet_strength=0.4, ) @@ -22,9 +24,11 @@ class TestImageGeneratorControlnet: reference_image_path="reference_controlnet_dev.png", output_image_path="output_controlnet_dev.png", controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME, - model_config=ModelConfig.dev(), + model_config=ModelConfig.dev_controlnet_canny(), steps=15, seed=42, + height=768, + width=493, prompt="The joker with a hat and a cane", controlnet_strength=0.4, ) @@ -34,11 +38,27 @@ class TestImageGeneratorControlnet: reference_image_path="reference_controlnet_dev_lora.png", output_image_path="output_controlnet_dev_lora.png", controlnet_image_path=TestImageGeneratorControlnet.CONTROLNET_REFERENCE_FILENAME, - model_config=ModelConfig.dev(), + model_config=ModelConfig.dev_controlnet_canny(), steps=15, seed=43, + height=768, + width=493, prompt="mkym this is made of wool, The joker with a hat and a cane", lora_paths=["FLUX-dev-lora-MiaoKa-Yarn-World.safetensors"], lora_scales=[1.0], controlnet_strength=0.4, ) + + def test_image_upscaling(self): + ImageGeneratorControlnetTestHelper.assert_matches_reference_image( + reference_image_path="reference_upscaled.png", + output_image_path="output_upscaler.png", + controlnet_image_path="low_res.jpg", + model_config=ModelConfig.dev_controlnet_upscaler(), + steps=20, + seed=42, + height=int(192 * 2), + width=int(320 * 2), + prompt="A man holding up a hand", + controlnet_strength=0.6, + ) diff --git a/tests/model_config/test_model_config.py b/tests/model_config/test_model_config.py index e18a8bb..4006899 100644 --- a/tests/model_config/test_model_config.py +++ b/tests/model_config/test_model_config.py @@ -64,6 +64,82 @@ def test_bfl_dev_fill_full_name(): assert model.requires_sigma_shift is True +def test_bfl_dev_depth(): + model = ModelConfig.from_name("dev-depth") + assert model.alias == "dev-depth" + assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev" + assert model.max_sequence_length == 512 + assert model.num_train_steps == 1000 + assert model.supports_guidance is True + assert model.requires_sigma_shift is True + + +def test_bfl_dev_depth_full_name(): + model = ModelConfig.from_name("black-forest-labs/FLUX.1-Depth-dev") + assert model.alias == "dev-depth" + assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev" + assert model.max_sequence_length == 512 + assert model.num_train_steps == 1000 + assert model.supports_guidance is True + assert model.requires_sigma_shift is True + + +def test_bfl_dev_redux(): + model = ModelConfig.from_name("dev-redux") + assert model.alias == "dev-redux" + assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev" + assert model.max_sequence_length == 512 + assert model.num_train_steps == 1000 + assert model.supports_guidance is True + assert model.requires_sigma_shift is True + + +def test_bfl_dev_redux_full_name(): + model = ModelConfig.from_name("black-forest-labs/FLUX.1-Redux-dev") + assert model.alias == "dev-redux" + assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev" + assert model.max_sequence_length == 512 + assert model.num_train_steps == 1000 + assert model.supports_guidance is True + assert model.requires_sigma_shift is True + + +def test_bfl_dev_controlnet_canny(): + model = ModelConfig.from_name("dev-controlnet-canny") + assert model.alias == "dev-controlnet-canny" + assert model.model_name == "black-forest-labs/FLUX.1-dev" + assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny" + assert model.max_sequence_length == 512 + assert model.num_train_steps == 1000 + assert model.supports_guidance is True + assert model.requires_sigma_shift is True + assert model.is_canny() is True + + +def test_bfl_schnell_controlnet_canny(): + model = ModelConfig.from_name("schnell-controlnet-canny") + assert model.alias == "schnell-controlnet-canny" + assert model.model_name == "black-forest-labs/FLUX.1-schnell" + assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny" + assert model.max_sequence_length == 256 + assert model.num_train_steps == 1000 + assert model.supports_guidance is False + assert model.requires_sigma_shift is False + assert model.is_canny() is True + + +def test_bfl_dev_controlnet_upscaler(): + model = ModelConfig.from_name("dev-controlnet-upscaler") + assert model.alias == "dev-controlnet-upscaler" + assert model.model_name == "black-forest-labs/FLUX.1-dev" + assert model.controlnet_model == "jasperai/Flux.1-dev-Controlnet-Upscaler" + assert model.max_sequence_length == 512 + assert model.num_train_steps == 1000 + assert model.supports_guidance is False + assert model.requires_sigma_shift is False + assert model.is_canny() is False + + def test_community_dev_fill_implicit_base_model(): model = ModelConfig.from_name("acme-lab/some-dev-fill-model") assert model.alias == "dev-fill" diff --git a/tests/resources/low_res.jpg b/tests/resources/low_res.jpg new file mode 100644 index 0000000..aca1bd6 Binary files /dev/null and b/tests/resources/low_res.jpg differ diff --git a/tests/resources/reference_depth_dev_from_image.png b/tests/resources/reference_depth_dev_from_image.png index 3e71373..98a7948 100644 Binary files a/tests/resources/reference_depth_dev_from_image.png and b/tests/resources/reference_depth_dev_from_image.png differ diff --git a/tests/resources/reference_upscaled.png b/tests/resources/reference_upscaled.png new file mode 100644 index 0000000..b094181 Binary files /dev/null and b/tests/resources/reference_upscaled.png differ