From dc081272d4d9beaf013d7dccafb684c3ba042940 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Sun, 30 Mar 2025 19:35:43 +0300 Subject: [PATCH 1/2] Update pre-commit tools --- .pre-commit-config.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b91cf35..440eef8 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,7 +1,7 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.6.9 + rev: v0.11.2 hooks: # Run the linter. - id: ruff @@ -11,6 +11,6 @@ repos: - id: ruff-format types_or: [python, pyi] - repo: https://github.com/crate-ci/typos - rev: v1.26.0 + rev: v1.31.0 hooks: - id: typos From 83a7bf3519e4585ab8dbbb5b67b356baa5afd3f5 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Sun, 30 Mar 2025 19:42:40 +0300 Subject: [PATCH 2/2] Clean up `fmt: off` comments * Use magic trailing commas instead of disabling formatting to keep args on separate lines * Scope `fmt: off`s better where that's not possible --- src/mflux/callbacks/callback.py | 18 +++++++---------- src/mflux/callbacks/callbacks.py | 14 ++++++------- src/mflux/callbacks/instances/canny_saver.py | 4 ++-- .../callbacks/instances/stepwise_handler.py | 14 ++++++------- .../in_context_lora/flux_in_context_lora.py | 8 ++++---- src/mflux/controlnet/flux_controlnet.py | 12 +++++------ .../controlnet/weight_handler_controlnet.py | 12 +++++------ src/mflux/dreambooth/dataset/iterator.py | 4 ++-- src/mflux/dreambooth/dreambooth.py | 20 +++++++++---------- .../dreambooth/dreambooth_initializer.py | 8 ++++---- .../dreambooth/lora_layers/lora_layers.py | 4 ++-- .../optimization/dreambooth_loss.py | 14 ++++++------- src/mflux/dreambooth/state/training_state.py | 2 +- src/mflux/flux/flux.py | 6 +++--- src/mflux/flux/flux_initializer.py | 12 +++++------ src/mflux/flux_tools/fill/flux_fill.py | 6 +++--- src/mflux/flux_tools/fill/mask_util.py | 4 ++-- src/mflux/latent_creator/latent_creator.py | 8 ++++---- .../transformer/common/attention_utils.py | 4 ++-- .../transformer/joint_transformer_block.py | 8 ++++---- .../transformer/single_transformer_block.py | 4 ++-- src/mflux/models/transformer/transformer.py | 2 +- .../models/vae/common/resnet_block_2d.py | 15 ++++++++------ src/mflux/models/vae/decoder/up_block_3.py | 4 +--- src/mflux/models/vae/decoder/up_block_4.py | 4 +--- src/mflux/models/vae/encoder/down_block_1.py | 4 +--- src/mflux/models/vae/encoder/down_block_2.py | 4 +--- src/mflux/models/vae/encoder/down_block_3.py | 4 +--- src/mflux/models/vae/encoder/down_block_4.py | 4 +--- src/mflux/post_processing/image_util.py | 4 ++-- src/mflux/train.py | 8 ++++---- src/mflux/ui/cli/parsers.py | 3 +-- src/mflux/weights/weight_handler_lora.py | 12 +++++------ tests/dreambooth/test_resume_training.py | 16 +++++++-------- .../dreambooth/test_train_and_load_weights.py | 10 +++++----- ...image_generation_in_context_test_helper.py | 4 ++-- .../helpers/image_generation_test_helper.py | 4 ++-- 37 files changed, 137 insertions(+), 151 deletions(-) diff --git a/src/mflux/callbacks/callback.py b/src/mflux/callbacks/callback.py index fe428ca..efd89e7 100644 --- a/src/mflux/callbacks/callback.py +++ b/src/mflux/callbacks/callback.py @@ -15,8 +15,7 @@ class BeforeLoopCallback(Protocol): latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, - ) -> None: # fmt: off - ... + ) -> None: ... class InLoopCallback(Protocol): @@ -27,9 +26,8 @@ class InLoopCallback(Protocol): prompt: str, latents: mx.array, config: RuntimeConfig, - time_steps: tqdm - ) -> None: # fmt: off - ... + time_steps: tqdm, + ) -> None: ... class AfterLoopCallback(Protocol): @@ -38,9 +36,8 @@ class AfterLoopCallback(Protocol): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig - ) -> None: # fmt: off - ... + config: RuntimeConfig, + ) -> None: ... class InterruptCallback(Protocol): @@ -51,6 +48,5 @@ class InterruptCallback(Protocol): prompt: str, latents: mx.array, config: RuntimeConfig, - time_steps: tqdm - ) -> None: # fmt: off - ... + time_steps: tqdm, + ) -> None: ... diff --git a/src/mflux/callbacks/callbacks.py b/src/mflux/callbacks/callbacks.py index f17af70..e415cec 100644 --- a/src/mflux/callbacks/callbacks.py +++ b/src/mflux/callbacks/callbacks.py @@ -14,7 +14,7 @@ class Callbacks: latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, - ): # fmt: off + ): for subscriber in CallbackRegistry.before_loop_callbacks(): subscriber.call_before_loop( seed=seed, @@ -31,8 +31,8 @@ class Callbacks: prompt: str, latents: mx.array, config: RuntimeConfig, - time_steps: tqdm - ): # fmt: off + time_steps: tqdm, + ): for subscriber in CallbackRegistry.in_loop_callbacks(): subscriber.call_in_loop( t=t, @@ -48,8 +48,8 @@ class Callbacks: seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig - ): # fmt: off + config: RuntimeConfig, + ): for subscriber in CallbackRegistry.after_loop_callbacks(): subscriber.call_after_loop(seed=seed, prompt=prompt, latents=latents, config=config) @@ -60,8 +60,8 @@ class Callbacks: prompt: str, latents: mx.array, config: RuntimeConfig, - time_steps: tqdm - ): # fmt: off + time_steps: tqdm, + ): for subscriber in CallbackRegistry.interrupt_callbacks(): subscriber.call_interrupt( t=t, diff --git a/src/mflux/callbacks/instances/canny_saver.py b/src/mflux/callbacks/instances/canny_saver.py index 548b291..4dcd361 100644 --- a/src/mflux/callbacks/instances/canny_saver.py +++ b/src/mflux/callbacks/instances/canny_saver.py @@ -24,5 +24,5 @@ class CannyImageSaver(BeforeLoopCallback): base, ext = os.path.splitext(self.path) ImageUtil.save_image( image=canny_image, - path=f"{base}_controlnet_canny{ext}" - ) # fmt: off + path=f"{base}_controlnet_canny{ext}", + ) diff --git a/src/mflux/callbacks/instances/stepwise_handler.py b/src/mflux/callbacks/instances/stepwise_handler.py index 8e0f0ba..7295ab6 100644 --- a/src/mflux/callbacks/instances/stepwise_handler.py +++ b/src/mflux/callbacks/instances/stepwise_handler.py @@ -30,7 +30,7 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): latents: mx.array, config: RuntimeConfig, canny_image: PIL.Image.Image | None = None, - ) -> None: # fmt: off + ) -> None: self._save_image( step=config.init_time_step, seed=seed, @@ -47,8 +47,8 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): prompt: str, latents: mx.array, config: RuntimeConfig, - time_steps: tqdm - ) -> None: # fmt: off + time_steps: tqdm, + ) -> None: self._save_image( step=t + 1, seed=seed, @@ -65,8 +65,8 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): prompt: str, latents: mx.array, config: RuntimeConfig, - time_steps: tqdm - ) -> None: # fmt: off + time_steps: tqdm, + ) -> None: self._save_composite(seed=seed) def _save_image( @@ -76,8 +76,8 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): prompt: str, latents: mx.array, config: RuntimeConfig, - time_steps: tqdm - ) -> None: # fmt: off + time_steps: tqdm, + ) -> None: unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) stepwise_decoded = self.flux.vae.decode(unpack_latents) generation_time = time_steps.format_dict["elapsed"] if time_steps is not None else 0 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 39ca300..90df0bc 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 @@ -145,7 +145,7 @@ class Flux1InContextLoRA(nn.Module): prompt=prompt, latents=latents, config=config, - ) # fmt: off + ) # 6. Decode the latent array and return the image latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) @@ -184,7 +184,7 @@ class Flux1InContextLoRA(nn.Module): latents: mx.array, encoded_image: mx.array, static_noise: mx.array, - ) -> mx.array: # fmt:off + ) -> mx.array: # 1. Unpack the latents unpacked = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) unpacked_static_noise = ArrayUtil.unpack_latents(latents=static_noise, height=config.height, width=config.width) @@ -196,8 +196,8 @@ class Flux1InContextLoRA(nn.Module): unpacked[:, :, :, 0:latent_width] = LatentCreator.add_noise_by_interpolation( clean=encoded_image[:, :, :, 0:latent_width], noise=unpacked_static_noise[:, :, :, 0:latent_width], - sigma=config.sigmas[t+1] - ) # fmt:off + sigma=config.sigmas[t + 1], + ) # 4. Repack the latents return ArrayUtil.pack_latents(latents=unpacked, height=config.height, width=config.width) diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 03b7401..addd8b9 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -71,8 +71,8 @@ class Flux1Controlnet(nn.Module): latents = LatentCreator.create( seed=seed, height=config.height, - width=config.width - ) # fmt: off + width=config.width, + ) # 3. Encode the prompt prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt( @@ -90,8 +90,8 @@ class Flux1Controlnet(nn.Module): prompt=prompt, latents=latents, config=config, - canny_image=canny_image - ) # fmt: off + canny_image=canny_image, + ) for t in time_steps: try: @@ -128,7 +128,7 @@ class Flux1Controlnet(nn.Module): latents=latents, config=config, time_steps=time_steps, - ) # fmt: off + ) # (Optional) Evaluate to enable progress tracking mx.eval(latents) @@ -150,7 +150,7 @@ class Flux1Controlnet(nn.Module): 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) diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py index f8b80c8..93988df 100644 --- a/src/mflux/controlnet/weight_handler_controlnet.py +++ b/src/mflux/controlnet/weight_handler_controlnet.py @@ -41,16 +41,16 @@ class WeightHandlerControlnet: return WeightHandlerControlnet( config=config, controlnet_transformer=weights, - meta_data=MetaData(quantization_level=quantization_level) - ) # fmt:off + meta_data=MetaData(quantization_level=quantization_level), + ) # Reshape and process the huggingface weights if "transformer_blocks" in weights: for block in weights["transformer_blocks"]: block["ff"] = { "linear1": block["ff"]["net"][0]["proj"], - "linear2": block["ff"]["net"][2] - } # fmt: off + "linear2": block["ff"]["net"][2], + } if block.get("ff_context") is not None: block["ff_context"] = { "linear1": block["ff_context"]["net"][0]["proj"], @@ -60,8 +60,8 @@ class WeightHandlerControlnet: return WeightHandlerControlnet( config=config, controlnet_transformer=weights, - meta_data=MetaData(quantization_level=quantization_level) - ) # fmt:off + meta_data=MetaData(quantization_level=quantization_level), + ) def num_transformer_blocks(self) -> int: return self.config["num_layers"] diff --git a/src/mflux/dreambooth/dataset/iterator.py b/src/mflux/dreambooth/dataset/iterator.py index 772b045..3eff4a8 100644 --- a/src/mflux/dreambooth/dataset/iterator.py +++ b/src/mflux/dreambooth/dataset/iterator.py @@ -143,8 +143,8 @@ class Iterator: data = ZipUtil.unzip( zip_path=training_spec.checkpoint_path, filename=iterator_path, - loader=lambda x: json.load(open(x, "r")) - ) # fmt: off + loader=lambda x: json.load(open(x, "r")), + ) return cls.from_dict(data, dataset) def get_validation_batch(self) -> Batch: diff --git a/src/mflux/dreambooth/dreambooth.py b/src/mflux/dreambooth/dreambooth.py index 16d8cf8..b6ad86e 100644 --- a/src/mflux/dreambooth/dreambooth.py +++ b/src/mflux/dreambooth/dreambooth.py @@ -13,23 +13,23 @@ from mflux.weights.weight_handler_lora import WeightHandlerLoRA class DreamBooth: @staticmethod def train( - flux: Flux1, - runtime_config: RuntimeConfig, - training_spec: TrainingSpec, - training_state: TrainingState - ): # fmt:off + flux: Flux1, + runtime_config: RuntimeConfig, + training_spec: TrainingSpec, + training_state: TrainingState, + ): # Freeze the model and assign the LoRA layers to the model flux.freeze() WeightHandlerLoRA.set_lora_layers( transformer_module=flux.transformer, - lora_layers=training_state.lora_layers - ) # fmt:off + lora_layers=training_state.lora_layers, + ) # Define loss computation as a function of a batch 'b' train_step_function = nn.value_and_grad( model=flux, - fn=lambda b: DreamBoothLoss.compute_loss(flux, runtime_config, b) - ) # fmt: off + fn=lambda b: DreamBoothLoss.compute_loss(flux, runtime_config, b), + ) # Setup progress bar batches = tqdm( @@ -59,7 +59,7 @@ class DreamBooth: seed=training_spec.seed, config=runtime_config.config, prompt=training_spec.instrumentation.validation_prompt, - ) # fmt: off + ) image.save(path=training_state.get_current_validation_image_path(training_spec)) del image flux.prompt_cache = {} diff --git a/src/mflux/dreambooth/dreambooth_initializer.py b/src/mflux/dreambooth/dreambooth_initializer.py index 0c85e90..8b11a51 100644 --- a/src/mflux/dreambooth/dreambooth_initializer.py +++ b/src/mflux/dreambooth/dreambooth_initializer.py @@ -21,8 +21,8 @@ class DreamBoothInitializer: # differently depending on if training starts from scratch or resumes from checkpoint. training_spec = TrainingSpec.resolve( config_path=config_path, - checkpoint_path=checkpoint_path - ) # fmt: off + checkpoint_path=checkpoint_path, + ) # Set global random seed to make training deterministic random.seed(training_spec.seed) @@ -58,8 +58,8 @@ class DreamBoothInitializer: ) iterator = Iterator.from_spec( training_spec=training_spec, - dataset=dataset - ) # fmt: off + dataset=dataset, + ) # Setup loss statistics statistics = Statistics.from_spec(training_spec=training_spec) diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py index edaa5cb..5ff3d9c 100644 --- a/src/mflux/dreambooth/lora_layers/lora_layers.py +++ b/src/mflux/dreambooth/lora_layers/lora_layers.py @@ -59,8 +59,8 @@ class LoRALayers: weights = WeightHandler( meta_data=MetaData(mflux_version=GeneratedImage.get_version()), - transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))['transformer'], - ) # fmt:off + transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))["transformer"], + ) return LoRALayers(weights=weights) diff --git a/src/mflux/dreambooth/optimization/dreambooth_loss.py b/src/mflux/dreambooth/optimization/dreambooth_loss.py index dbc23e7..1740cfa 100644 --- a/src/mflux/dreambooth/optimization/dreambooth_loss.py +++ b/src/mflux/dreambooth/optimization/dreambooth_loss.py @@ -31,9 +31,9 @@ class DreamBoothLoss: low=0, high=config.num_inference_steps, shape=[], - key=mx.random.key(time_seed) - ) - ) # fmt: off + key=mx.random.key(time_seed), + ), + ) # Get the clean image latent clean_image = example.clean_latents @@ -42,15 +42,15 @@ class DreamBoothLoss: pure_noise = mx.random.normal( shape=clean_image.shape, dtype=Config.precision, - key=mx.random.key(noise_seed) - ) # fmt: off + key=mx.random.key(noise_seed), + ) # By linear interpolation between the clean image and pure noise, construct a latent at time t latents_t = LatentCreator.add_noise_by_interpolation( clean=clean_image, noise=pure_noise, - sigma=config.sigmas[t] - ) # fmt: off + sigma=config.sigmas[t], + ) # Predict the noise from timestep t predicted_noise = flux.transformer( diff --git a/src/mflux/dreambooth/state/training_state.py b/src/mflux/dreambooth/state/training_state.py index a9102de..33d61ef 100644 --- a/src/mflux/dreambooth/state/training_state.py +++ b/src/mflux/dreambooth/state/training_state.py @@ -49,7 +49,7 @@ class TrainingState: iterator_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_ITERATOR}.json" loss_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LOSS_FILE}.json" config_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_CONFIG_FILE}.json" - checkpoint_path = Path(temp_dir) / f"{DREAMBOOTH_FILE_NAME_CHECKPOINT}.json" # fmt: off + checkpoint_path = Path(temp_dir) / f"{DREAMBOOTH_FILE_NAME_CHECKPOINT}.json" paths = [optimizer_path, lora_path, iterator_path, loss_path, config_path, checkpoint_path] # Save individual files to temporary directory diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index f9aabbe..c43a63d 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -83,7 +83,7 @@ class Flux1(nn.Module): prompt=prompt, latents=latents, config=config, - ) # fmt: off + ) for t in time_steps: try: @@ -108,7 +108,7 @@ class Flux1(nn.Module): latents=latents, config=config, time_steps=time_steps, - ) # fmt: off + ) # (Optional) Evaluate to enable progress tracking mx.eval(latents) @@ -130,7 +130,7 @@ class Flux1(nn.Module): 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) diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index 0ecab05..df04f2a 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/src/mflux/flux/flux_initializer.py @@ -41,8 +41,8 @@ class FluxInitializer: ) flux_model.t5_tokenizer = TokenizerT5( tokenizer=tokenizers.t5, - max_length=model_config.max_sequence_length - ) # fmt: off + max_length=model_config.max_sequence_length, + ) flux_model.clip_tokenizer = TokenizerCLIP( tokenizer=tokenizers.clip, ) @@ -50,8 +50,8 @@ class FluxInitializer: # 2. Load the regular weights weights = WeightHandler.load_regular_weights( repo_id=model_config.model_name, - local_path=local_path - ) # fmt: off + local_path=local_path, + ) # 3. Initialize all models flux_model.vae = VAE() @@ -85,8 +85,8 @@ class FluxInitializer: ) WeightHandlerLoRA.set_lora_weights( transformer=flux_model.transformer, - loras=lora_weights - ) # fmt: off + loras=lora_weights, + ) @staticmethod def init_controlnet( diff --git a/src/mflux/flux_tools/fill/flux_fill.py b/src/mflux/flux_tools/fill/flux_fill.py index ac80547..503529a 100644 --- a/src/mflux/flux_tools/fill/flux_fill.py +++ b/src/mflux/flux_tools/fill/flux_fill.py @@ -87,7 +87,7 @@ class Flux1Fill(nn.Module): prompt=prompt, latents=latents, config=config, - ) # fmt: off + ) for t in time_steps: try: @@ -115,7 +115,7 @@ class Flux1Fill(nn.Module): latents=latents, config=config, time_steps=time_steps, - ) # fmt: off + ) # (Optional) Evaluate to enable progress tracking mx.eval(latents) @@ -137,7 +137,7 @@ class Flux1Fill(nn.Module): 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) diff --git a/src/mflux/flux_tools/fill/mask_util.py b/src/mflux/flux_tools/fill/mask_util.py index 72406be..d2ef628 100644 --- a/src/mflux/flux_tools/fill/mask_util.py +++ b/src/mflux/flux_tools/fill/mask_util.py @@ -13,8 +13,8 @@ class MaskUtil: config: RuntimeConfig, latents: mx.array, img_path: str, - mask_path: str | None - ) -> mx.array: # fmt: off + mask_path: str | None, + ) -> mx.array: if not img_path or not mask_path: # Return empty latents if no image or mask is provided return mx.zeros((1, 0, 0)) diff --git a/src/mflux/latent_creator/latent_creator.py b/src/mflux/latent_creator/latent_creator.py index aff3668..046ee4f 100644 --- a/src/mflux/latent_creator/latent_creator.py +++ b/src/mflux/latent_creator/latent_creator.py @@ -28,8 +28,8 @@ class LatentCreator: ) -> mx.array: return mx.random.normal( shape=[1, (height // 16) * (width // 16), 64], - key=mx.random.key(seed) - ) # fmt: off + key=mx.random.key(seed), + ) @staticmethod def create_for_txt2img_or_img2img( @@ -67,8 +67,8 @@ class LatentCreator: return LatentCreator.add_noise_by_interpolation( clean=latents, noise=pure_noise, - sigma=sigma - ) # fmt: off + sigma=sigma, + ) @staticmethod def encode_image(height: int, width: int, img2img: Img2Img): diff --git a/src/mflux/models/transformer/common/attention_utils.py b/src/mflux/models/transformer/common/attention_utils.py index 424ae48..7674424 100644 --- a/src/mflux/models/transformer/common/attention_utils.py +++ b/src/mflux/models/transformer/common/attention_utils.py @@ -37,8 +37,8 @@ class AttentionUtils: value: mx.array, batch_size: int, num_heads: int, - head_dim: int - ) -> mx.array: # fmt: off + head_dim: int, + ) -> mx.array: scale = 1 / mx.sqrt(query.shape[-1]) hidden_states = scaled_dot_product_attention(query, key, value, scale=scale) diff --git a/src/mflux/models/transformer/joint_transformer_block.py b/src/mflux/models/transformer/joint_transformer_block.py index 2eb5419..626372e 100644 --- a/src/mflux/models/transformer/joint_transformer_block.py +++ b/src/mflux/models/transformer/joint_transformer_block.py @@ -28,14 +28,14 @@ class JointTransformerBlock(nn.Module): # 1a. Compute norm for hidden_states norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( hidden_states=hidden_states, - text_embeddings=text_embeddings - ) # fmt: off + text_embeddings=text_embeddings, + ) # 1b. Compute norm for encoder_hidden_states norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( hidden_states=encoder_hidden_states, - text_embeddings=text_embeddings - ) # fmt: off + text_embeddings=text_embeddings, + ) # 2. Compute attention attn_output, context_attn_output = self.attn( diff --git a/src/mflux/models/transformer/single_transformer_block.py b/src/mflux/models/transformer/single_transformer_block.py index a2344ce..a188b51 100644 --- a/src/mflux/models/transformer/single_transformer_block.py +++ b/src/mflux/models/transformer/single_transformer_block.py @@ -28,8 +28,8 @@ class SingleTransformerBlock(nn.Module): # 1. Compute norm for hidden_states norm_hidden_states, gate = self.norm( hidden_states=hidden_states, - text_embeddings=text_embeddings - ) # fmt: off + text_embeddings=text_embeddings, + ) # 2. Compute attention attn_output = self.attn( diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py index 0a20375..84e84b0 100644 --- a/src/mflux/models/transformer/transformer.py +++ b/src/mflux/models/transformer/transformer.py @@ -172,7 +172,7 @@ class Transformer(nn.Module): idx: int, blocks: mx.array, controlnet_samples: list[mx.array] | None, - ) -> mx.array | None: # fmt: off + ) -> mx.array | None: if controlnet_samples is None: return None diff --git a/src/mflux/models/vae/common/resnet_block_2d.py b/src/mflux/models/vae/common/resnet_block_2d.py index f235944..356a627 100644 --- a/src/mflux/models/vae/common/resnet_block_2d.py +++ b/src/mflux/models/vae/common/resnet_block_2d.py @@ -4,7 +4,6 @@ from mlx import nn from mflux.config.config import Config -# fmt: off class ResnetBlock2D(nn.Module): def __init__( self, @@ -48,11 +47,15 @@ class ResnetBlock2D(nn.Module): padding=(1, 1), ) self.is_conv_shortcut = is_conv_shortcut - self.conv_shortcut = None if not is_conv_shortcut else nn.Conv2d( - in_channels=conv_shortcut_in, - out_channels=conv_shortcut_out, - kernel_size=(1, 1), - stride=(1, 1), + self.conv_shortcut = ( + nn.Conv2d( + in_channels=conv_shortcut_in, + out_channels=conv_shortcut_out, + kernel_size=(1, 1), + stride=(1, 1), + ) + if is_conv_shortcut + else None ) def __call__(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/decoder/up_block_3.py b/src/mflux/models/vae/decoder/up_block_3.py index c3f8a46..b7c45ff 100644 --- a/src/mflux/models/vae/decoder/up_block_3.py +++ b/src/mflux/models/vae/decoder/up_block_3.py @@ -8,13 +8,11 @@ from mflux.models.vae.decoder.up_sampler import UpSampler class UpBlock3(nn.Module): def __init__(self): super().__init__() - # fmt: off self.resnets = [ ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256, is_conv_shortcut=True, conv_shortcut_in=512, conv_shortcut_out=256), ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256), ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256), - ] - # fmt: on + ] # fmt: off self.upsamplers = [UpSampler(conv_in=256, conv_out=256)] def __call__(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/decoder/up_block_4.py b/src/mflux/models/vae/decoder/up_block_4.py index b69235c..4d9b6a9 100644 --- a/src/mflux/models/vae/decoder/up_block_4.py +++ b/src/mflux/models/vae/decoder/up_block_4.py @@ -7,13 +7,11 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D class UpBlock4(nn.Module): def __init__(self): super().__init__() - # fmt: off self.resnets = [ ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128, is_conv_shortcut=True, conv_shortcut_in=256, conv_shortcut_out=128), ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), - ] - # fmt: off + ] # fmt: off def __call__(self, input_array: mx.array) -> mx.array: hidden_states = self.resnets[0](input_array) diff --git a/src/mflux/models/vae/encoder/down_block_1.py b/src/mflux/models/vae/encoder/down_block_1.py index 96b1ee8..2abd5a8 100644 --- a/src/mflux/models/vae/encoder/down_block_1.py +++ b/src/mflux/models/vae/encoder/down_block_1.py @@ -9,11 +9,9 @@ class DownBlock1(nn.Module): def __init__(self): super().__init__() self.resnets = [ - # fmt: off ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), - # fmt: on - ] + ] # fmt: off self.downsamplers = [DownSampler(conv_in=128, conv_out=128)] def __call__(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/encoder/down_block_2.py b/src/mflux/models/vae/encoder/down_block_2.py index cf258e1..88beddc 100644 --- a/src/mflux/models/vae/encoder/down_block_2.py +++ b/src/mflux/models/vae/encoder/down_block_2.py @@ -8,12 +8,10 @@ from mflux.models.vae.encoder.down_sampler import DownSampler class DownBlock2(nn.Module): def __init__(self): super().__init__() - # fmt: off self.resnets = [ ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256, is_conv_shortcut=True, conv_shortcut_in=128, conv_shortcut_out=256), ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256), - ] - # fmt: on + ] # fmt: off self.downsamplers = [DownSampler(conv_in=256, conv_out=256)] def __call__(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/encoder/down_block_3.py b/src/mflux/models/vae/encoder/down_block_3.py index f94a5f9..80cca20 100644 --- a/src/mflux/models/vae/encoder/down_block_3.py +++ b/src/mflux/models/vae/encoder/down_block_3.py @@ -8,12 +8,10 @@ from mflux.models.vae.encoder.down_sampler import DownSampler class DownBlock3(nn.Module): def __init__(self): super().__init__() - # fmt: off self.resnets = [ ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512, is_conv_shortcut=True, conv_shortcut_in=256, conv_shortcut_out=512), ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), - ] - # fmt: on + ] # fmt: off self.downsamplers = [DownSampler(conv_in=512, conv_out=512)] def __call__(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/encoder/down_block_4.py b/src/mflux/models/vae/encoder/down_block_4.py index 7ffd825..95c9bf5 100644 --- a/src/mflux/models/vae/encoder/down_block_4.py +++ b/src/mflux/models/vae/encoder/down_block_4.py @@ -7,12 +7,10 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D class DownBlock4(nn.Module): def __init__(self): super().__init__() - # fmt: off self.resnets = [ ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), - ] - # fmt: on + ] # fmt: off def __call__(self, input_array: mx.array) -> mx.array: hidden_states = self.resnets[0](input_array) diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 08af878..bf65bdf 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -203,8 +203,8 @@ class ImageUtil: path: t.Union[str, pathlib.Path], metadata: dict | None = None, export_json_metadata: bool = False, - overwrite: bool = False - ) -> None: # fmt: off + overwrite: bool = False, + ) -> None: file_path = pathlib.Path(path) file_path.parent.mkdir(parents=True, exist_ok=True) file_name = file_path.stem diff --git a/src/mflux/train.py b/src/mflux/train.py index a8f07ae..9501373 100644 --- a/src/mflux/train.py +++ b/src/mflux/train.py @@ -13,16 +13,16 @@ def main(): flux, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( config_path=args.train_config, - checkpoint_path=args.train_checkpoint - ) # fmt: off + checkpoint_path=args.train_checkpoint, + ) try: DreamBooth.train( flux=flux, runtime_config=runtime_config, training_spec=training_spec, - training_state=training_state - ) # fmt: off + training_state=training_state, + ) except StopTrainingException as stop_exc: training_state.save(training_spec) print(stop_exc) diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index f38dbae..9353443 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -54,7 +54,7 @@ class CommandLineParser(argparse.ArgumentParser): self.add_argument("--base-model", type=str, required=False, choices=ui_defaults.MODEL_CHOICES, help="When using a third-party huggingface model, explicitly specify whether the base model is dev or schnell") 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_lora_arguments(self) -> None: # fmt: off + def add_lora_arguments(self) -> None: self.supports_lora = True lora_group = self.add_argument_group("LoRA configuration") lora_group.add_argument("--lora-style", type=str, choices=sorted(LORA_NAME_MAP.keys()), help="Style of the LoRA to use (e.g., 'storyboard' for film storyboard style)") @@ -62,7 +62,6 @@ class CommandLineParser(argparse.ArgumentParser): self.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") lora_group.add_argument("--lora-name", type=str, help="Name of the LoRA to download from Hugging Face") lora_group.add_argument("--lora-repo-id", type=str, default=LORA_REPO_ID, help=f"Hugging Face repository ID for LoRAs (default: {LORA_REPO_ID})") - # fmt: on def _add_image_generator_common_arguments(self) -> None: self.supports_image_generation = True diff --git a/src/mflux/weights/weight_handler_lora.py b/src/mflux/weights/weight_handler_lora.py index da514f6..740873a 100644 --- a/src/mflux/weights/weight_handler_lora.py +++ b/src/mflux/weights/weight_handler_lora.py @@ -60,8 +60,8 @@ class WeightHandlerLoRA: WeightHandlerLoRA.set_lora_layers( transformer_module=transformer, - lora_layers=LoRALayers(weights=fused_weights) - ) # fmt:off + lora_layers=LoRALayers(weights=fused_weights), + ) @staticmethod def _fuse_multiple_lora_dicts(dicts: list[dict]) -> dict: @@ -143,13 +143,13 @@ class WeightHandlerLoRA: for i, weights in enumerate(transformer_blocks): LoRALayers.set_transformer_block( transformer_block=transformer_module.transformer_blocks[i], - dictionary=weights - ) # fmt:off + dictionary=weights, + ) # Handle single_transformer_blocks single_transformer_blocks = transformer.get("single_transformer_blocks", []) for i, weights in enumerate(single_transformer_blocks): LoRALayers.set_single_transformer_block( single_transformer_block=transformer_module.single_transformer_blocks[i], - dictionary=weights - ) # fmt:off + dictionary=weights, + ) diff --git a/tests/dreambooth/test_resume_training.py b/tests/dreambooth/test_resume_training.py index 1d3c4c8..7423372 100644 --- a/tests/dreambooth/test_resume_training.py +++ b/tests/dreambooth/test_resume_training.py @@ -21,14 +21,14 @@ class TestResumeTraining: # Given: A small training run from scratch for 5 steps (as described in the config)... fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( config_path="tests/dreambooth/config/train.json", - checkpoint_path=None - ) # fmt:off + checkpoint_path=None, + ) DreamBooth.train( flux=fluxA, runtime_config=runtime_config, training_spec=training_spec, - training_state=training_state - ) # fmt: off + training_state=training_state, + ) del fluxA, runtime_config, training_spec, training_state # ...where we can inspect the training state after 5 runs... adapter_after_5_steps = ZipUtil.unzip( @@ -43,14 +43,14 @@ class TestResumeTraining: # When: Resuming the training from step 3... fluxB, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( config_path=None, - checkpoint_path=CHECKPOINT_3 - ) # fmt:off + checkpoint_path=CHECKPOINT_3, + ) DreamBooth.train( flux=fluxB, runtime_config=runtime_config, training_spec=training_spec, - training_state=training_state - ) # fmt: off + training_state=training_state, + ) del fluxB, runtime_config, training_spec, training_state # ...where we can inspect the training state after 2 additional runs... adapter_after_5_steps_resumed = ZipUtil.unzip( diff --git a/tests/dreambooth/test_train_and_load_weights.py b/tests/dreambooth/test_train_and_load_weights.py index bb6456e..0d0697e 100644 --- a/tests/dreambooth/test_train_and_load_weights.py +++ b/tests/dreambooth/test_train_and_load_weights.py @@ -22,14 +22,14 @@ class TestTrainAndLoadWeights: # Given: A small training run from scratch for 5 steps (as described in the config)... fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( config_path="tests/dreambooth/config/train.json", - checkpoint_path=None - ) # fmt:off + checkpoint_path=None, + ) DreamBooth.train( flux=fluxA, runtime_config=runtime_config, training_spec=training_spec, - training_state=training_state - ) # fmt: off + training_state=training_state, + ) # ...we generate an image with the flux instance with the trained weights image1 = fluxA.generate_image( seed=42, @@ -50,7 +50,7 @@ class TestTrainAndLoadWeights: quantize=4, lora_paths=[LORA_FILE], lora_scales=[1.0], - ) # fmt: off + ) # ...and generating the same image from that image2 = fluxB.generate_image( diff --git a/tests/image_generation/helpers/image_generation_in_context_test_helper.py b/tests/image_generation/helpers/image_generation_in_context_test_helper.py index ed8e4d3..08712bd 100644 --- a/tests/image_generation/helpers/image_generation_in_context_test_helper.py +++ b/tests/image_generation/helpers/image_generation_in_context_test_helper.py @@ -41,8 +41,8 @@ class ImageGeneratorInContextTestHelper: lora_names=[get_lora_filename(lora_style)] if lora_style else None, lora_repo_id=LORA_REPO_ID if lora_style else None, lora_paths=lora_paths, - lora_scales=lora_scales - ) # fmt: off + lora_scales=lora_scales, + ) # when image = flux.generate_image( seed=seed, diff --git a/tests/image_generation/helpers/image_generation_test_helper.py b/tests/image_generation/helpers/image_generation_test_helper.py index f816e44..6af307c 100644 --- a/tests/image_generation/helpers/image_generation_test_helper.py +++ b/tests/image_generation/helpers/image_generation_test_helper.py @@ -34,8 +34,8 @@ class ImageGeneratorTestHelper: model_config=model_config, quantize=8, lora_paths=lora_paths, - lora_scales=lora_scales - ) # fmt: off + lora_scales=lora_scales, + ) # when image = flux.generate_image( seed=seed,