monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clean up fmt: off comments

Browse Source 
* Use magic trailing commas instead of disabling formatting to keep args on separate lines
* Scope `fmt: off`s better where that's not possible
This commit is contained in:
Aarni Koskela 2025-03-30 19:42:40 +03:00 committed by filipstrand
parent dc081272d4
commit 83a7bf3519
37 changed files with 137 additions and 151 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
src/mflux/callbacks/callback.py
View File

@ -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: ...

14
src/mflux/callbacks/callbacks.py
View File

@ -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,

4
src/mflux/callbacks/instances/canny_saver.py
View File

@ -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}",
)

14
src/mflux/callbacks/instances/stepwise_handler.py
View File

@ -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

8
src/mflux/community/in_context_lora/flux_in_context_lora.py
View File

@ -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)

12
src/mflux/controlnet/flux_controlnet.py
View File

@ -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)

12
src/mflux/controlnet/weight_handler_controlnet.py
View File

@ -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"]

4
src/mflux/dreambooth/dataset/iterator.py
View File

@ -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:

20
src/mflux/dreambooth/dreambooth.py
View File

@ -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 = {}

8
src/mflux/dreambooth/dreambooth_initializer.py
View File

@ -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)

4
src/mflux/dreambooth/lora_layers/lora_layers.py
View File

@ -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)

14
src/mflux/dreambooth/optimization/dreambooth_loss.py
View File

@ -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(

2
src/mflux/dreambooth/state/training_state.py
View File

@ -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

6
src/mflux/flux/flux.py
View File

@ -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)

12
src/mflux/flux/flux_initializer.py
View File

@ -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(

6
src/mflux/flux_tools/fill/flux_fill.py
View File

@ -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)

4
src/mflux/flux_tools/fill/mask_util.py
View File

@ -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))

8
src/mflux/latent_creator/latent_creator.py
View File

@ -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):

4
src/mflux/models/transformer/common/attention_utils.py
View File

@ -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)

8
src/mflux/models/transformer/joint_transformer_block.py
View File

@ -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(

4
src/mflux/models/transformer/single_transformer_block.py
View File

@ -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(

2
src/mflux/models/transformer/transformer.py
View File

@ -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

15
src/mflux/models/vae/common/resnet_block_2d.py
View File

@ -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:

4
src/mflux/models/vae/decoder/up_block_3.py
View File

@ -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:

4
src/mflux/models/vae/decoder/up_block_4.py
View File

@ -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)

4
src/mflux/models/vae/encoder/down_block_1.py
View File

@ -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:

4
src/mflux/models/vae/encoder/down_block_2.py
View File

@ -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:

4
src/mflux/models/vae/encoder/down_block_3.py
View File

@ -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:

4
src/mflux/models/vae/encoder/down_block_4.py
View File

@ -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)

4
src/mflux/post_processing/image_util.py
View File

@ -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

8
src/mflux/train.py
View File

@ -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)

3
src/mflux/ui/cli/parsers.py
View File

@ -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

12
src/mflux/weights/weight_handler_lora.py
View File

@ -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,
)

16
tests/dreambooth/test_resume_training.py
View File

@ -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(

10
tests/dreambooth/test_train_and_load_weights.py
View File

@ -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(

4
tests/image_generation/helpers/image_generation_in_context_test_helper.py
View File

@ -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,

4
tests/image_generation/helpers/image_generation_test_helper.py
View File

@ -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,