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

@@ -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
+        time_steps: tqdm,
-    ) -> None:  # fmt: off
+    ) -> None: ...
-        ...
 
 
 class AfterLoopCallback(Protocol):
@@ -38,9 +36,8 @@ class AfterLoopCallback(Protocol):
         seed: int,
         prompt: str,
         latents: mx.array,
-        config: RuntimeConfig
+        config: RuntimeConfig,
-    ) -> None:  # fmt: off
+    ) -> None: ...
-        ...
 
 
 class InterruptCallback(Protocol):
@@ -51,6 +48,5 @@ class InterruptCallback(Protocol):
         prompt: str,
         latents: mx.array,
         config: RuntimeConfig,
-        time_steps: tqdm
+        time_steps: tqdm,
-    ) -> None:  # fmt: off
+    ) -> None: ...
-        ...

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
+        time_steps: tqdm,
-    ):  # fmt: off
+    ):
         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
+        config: RuntimeConfig,
-    ):  # fmt: off
+    ):
         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
+        time_steps: tqdm,
-    ):  # fmt: off
+    ):
         for subscriber in CallbackRegistry.interrupt_callbacks():
             subscriber.call_interrupt(
                 t=t,

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}"
+            path=f"{base}_controlnet_canny{ext}",
-        )  # fmt: off
+        )

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
+        time_steps: tqdm,
-    ) -> None:  # fmt: off
+    ) -> 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
+        time_steps: tqdm,
-    ) -> None:  # fmt: off
+    ) -> 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
+        time_steps: tqdm,
-    ) -> None:  # fmt: off
+    ) -> 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

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]
+            sigma=config.sigmas[t + 1],
-        )  # fmt:off
+        )
 
         # 4. Repack the latents
         return ArrayUtil.pack_latents(latents=unpacked, height=config.height, width=config.width)

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
+            width=config.width,
-        )  # fmt: off
+        )
 
         # 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
+            canny_image=canny_image,
-        )  # fmt: off
+        )
 
         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)

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)
+                meta_data=MetaData(quantization_level=quantization_level),
-            )  # fmt:off
+            )
 
         # 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]
+                    "linear2": block["ff"]["net"][2],
-                }  # fmt: off
+                }
                 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)
+            meta_data=MetaData(quantization_level=quantization_level),
-        )  # fmt:off
+        )
 
     def num_transformer_blocks(self) -> int:
         return self.config["num_layers"]

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"))
+            loader=lambda x: json.load(open(x, "r")),
-        )  # fmt: off
+        )
         return cls.from_dict(data, dataset)
 
     def get_validation_batch(self) -> Batch:

src/mflux/dreambooth/dreambooth.py

@@ -13,23 +13,23 @@ from mflux.weights.weight_handler_lora import WeightHandlerLoRA
 class DreamBooth:
     @staticmethod
     def train(
-            flux: Flux1,
+        flux: Flux1,
-            runtime_config: RuntimeConfig,
+        runtime_config: RuntimeConfig,
-            training_spec: TrainingSpec,
+        training_spec: TrainingSpec,
-            training_state: TrainingState
+        training_state: TrainingState,
-    ):  # fmt:off
+    ):
         # 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
+            lora_layers=training_state.lora_layers,
-        )  # fmt:off
+        )
 
         # 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)
+            fn=lambda b: DreamBoothLoss.compute_loss(flux, runtime_config, b),
-        )  # fmt: off
+        )
 
         # 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 = {}

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
+            checkpoint_path=checkpoint_path,
-        )  # fmt: off
+        )
 
         # 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
+            dataset=dataset,
-        )  # fmt: off
+        )
 
         # Setup loss statistics
         statistics = Statistics.from_spec(training_spec=training_spec)

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'],
+                transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))["transformer"],
-            )  # fmt:off
+            )
 
             return LoRALayers(weights=weights)
 

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)
+                key=mx.random.key(time_seed),
-            )
+            ),
-        )  # fmt: off
+        )
 
         # 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)
+            key=mx.random.key(noise_seed),
-        )  # fmt: off
+        )
 
         # 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]
+            sigma=config.sigmas[t],
-        )  # fmt: off
+        )
 
         # Predict the noise from timestep t
         predicted_noise = flux.transformer(

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

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)

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
+            max_length=model_config.max_sequence_length,
-        )  # fmt: off
+        )
         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
+            local_path=local_path,
-        )  # fmt: off
+        )
 
         # 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
+            loras=lora_weights,
-        )  # fmt: off
+        )
 
     @staticmethod
     def init_controlnet(

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)

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
+        mask_path: str | None,
-    ) -> mx.array:  # fmt: off
+    ) -> 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))

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)
+            key=mx.random.key(seed),
-        )  # fmt: off
+        )
 
     @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
+                sigma=sigma,
-            )  # fmt: off
+            )
 
     @staticmethod
     def encode_image(height: int, width: int, img2img: Img2Img):

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
+        head_dim: int,
-    ) -> mx.array:  # fmt: off
+    ) -> mx.array:
         scale = 1 / mx.sqrt(query.shape[-1])
         hidden_states = scaled_dot_product_attention(query, key, value, scale=scale)
 

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
+            text_embeddings=text_embeddings,
-        )  # fmt: off
+        )
 
         # 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
+            text_embeddings=text_embeddings,
-        )  # fmt: off
+        )
 
         # 2. Compute attention
         attn_output, context_attn_output = self.attn(

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
+            text_embeddings=text_embeddings,
-        )  # fmt: off
+        )
 
         # 2. Compute attention
         attn_output = self.attn(

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
 

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(
+        self.conv_shortcut = (
-            in_channels=conv_shortcut_in,
+            nn.Conv2d(
-            out_channels=conv_shortcut_out,
+                in_channels=conv_shortcut_in,
-            kernel_size=(1, 1),
+                out_channels=conv_shortcut_out,
-            stride=(1, 1),
+                kernel_size=(1, 1),
+                stride=(1, 1),
+            )
+            if is_conv_shortcut
+            else None
         )
 
     def __call__(self, input_array: mx.array) -> mx.array:

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: off
-        # fmt: on
         self.upsamplers = [UpSampler(conv_in=256, conv_out=256)]
 
     def __call__(self, input_array: mx.array) -> mx.array:

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)

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:

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: off
-        # fmt: on
         self.downsamplers = [DownSampler(conv_in=256, conv_out=256)]
 
     def __call__(self, input_array: mx.array) -> mx.array:

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: off
-        # fmt: on
         self.downsamplers = [DownSampler(conv_in=512, conv_out=512)]
 
     def __call__(self, input_array: mx.array) -> mx.array:

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: off
-        # fmt: on
 
     def __call__(self, input_array: mx.array) -> mx.array:
         hidden_states = self.resnets[0](input_array)

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
+        overwrite: bool = False,
-    ) -> None:  # fmt: off
+    ) -> None:
         file_path = pathlib.Path(path)
         file_path.parent.mkdir(parents=True, exist_ok=True)
         file_name = file_path.stem

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
+        checkpoint_path=args.train_checkpoint,
-    )  # fmt: off
+    )
 
     try:
         DreamBooth.train(
             flux=flux,
             runtime_config=runtime_config,
             training_spec=training_spec,
-            training_state=training_state
+            training_state=training_state,
-        )  # fmt: off
+        )
     except StopTrainingException as stop_exc:
         training_state.save(training_spec)
         print(stop_exc)

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

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)
+                lora_layers=LoRALayers(weights=fused_weights),
-            )  # fmt:off
+            )
 
     @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
+                dictionary=weights,
-            )  # fmt:off
+            )
 
         # 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
+                dictionary=weights,
-            )  # fmt:off
+            )

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
+                checkpoint_path=None,
-            )  # fmt:off
+            )
             DreamBooth.train(
                 flux=fluxA,
                 runtime_config=runtime_config,
                 training_spec=training_spec,
-                training_state=training_state
+                training_state=training_state,
-            )  # fmt: off
+            )
             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
+                checkpoint_path=CHECKPOINT_3,
-            )  # fmt:off
+            )
             DreamBooth.train(
                 flux=fluxB,
                 runtime_config=runtime_config,
                 training_spec=training_spec,
-                training_state=training_state
+                training_state=training_state,
-            )  # fmt: off
+            )
             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(

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
+                checkpoint_path=None,
-            )  # fmt:off
+            )
             DreamBooth.train(
                 flux=fluxA,
                 runtime_config=runtime_config,
                 training_spec=training_spec,
-                training_state=training_state
+                training_state=training_state,
-            )  # fmt: off
+            )
             # ...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(

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
+                lora_scales=lora_scales,
-            )  # fmt: off
+            )
             # when
             image = flux.generate_image(
                 seed=seed,

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
+                lora_scales=lora_scales,
-            )  # fmt: off
+            )
             # when
             image = flux.generate_image(
                 seed=seed,