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

Some fixes: Removed time in plotting, fix off-by-one error in training spec, update y-axis on loss plot

Browse Source
This commit is contained in:
filipstrand 2025-01-27 22:39:29 +01:00
parent 401da9fd75
commit f94edc4c3b
4 changed files with 53 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
README.md
View File

@ -159,7 +159,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
- **`--seed`** (optional, repeatable `int` args, default: `None`): 1 or more seeds for random number generation. e.g. `--seed 42` or `--seed 123 456 789`. Default is a single time-based value.
- **`--auto-seeds`** (optional, `int`, default: `None`): Auto generate N random Seeds in a series of image generations. Superceded by `--seed` arg and `seed` values in `--config-from-metadata` files.
- **`--auto-seeds`** (optional, `int`, default: `None`): Auto generate N random Seeds in a series of image generations. Superseded by `--seed` arg and `seed` values in `--config-from-metadata` files.
- **`--height`** (optional, `int`, default: `1024`): Height of the output image in pixels.
@ -689,6 +689,28 @@ The maximum range available for the different layer categories are:
- `start: 0`
- `end: 38`
<details>
<summary>Specify individual layers</summary>
For even more precision, you can specify individual block indices to train like so:
```json
"lora_layers": {
"single_transformer_blocks": {
"block_range": {
"indices": [
0,
1,
7,
19,
20
],
...
},
...
```
</details>
*⚠️ Note: As the joint transformer blocks (`transformer_blocks`) - are placed earlier on in the sequence of computations, they will require more resources to train.
In other words, training later layers, such as only the `single_transformer_blocks` should be faster. However, training too few / only later layers might result in a faster but unsuccessful training.*

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

@ -74,7 +74,7 @@ class LoRALayers:
lora_layers = {}
for idx in block_indices:
if idx >= len(blocks):
raise IndexError(f"Indice {idx} over range")
raise IndexError(f"Index {idx} over range")
block = blocks[idx]
for layer_type in block_spec.layer_types:

12
src/mflux/dreambooth/state/training_spec.py
View File

@ -3,7 +3,7 @@ import json
import os
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import List, Optional
from typing import List
from mflux.dreambooth.state.zip_util import ZipUtil
@ -56,16 +56,16 @@ class StatisticsSpec:
@dataclass
class BlockRange:
start: Optional[int] = None
end: Optional[int] = None
indices: Optional[List[int]] = None
start: int | None = None
end: int | None = None
indices: List[int] | None = None
def get_blocks(self) -> List[int]:
if self.indices:
return self.indices
if self.start is not None and self.end is not None:
return list(range(self.start, self.end - 1))
raise ValueError("Either ‘start’ and ‘end’ or ‘indices’ must be provided.")
return list(range(self.start, self.end))
raise ValueError("Either 'start' and 'end' or 'indices' must be provided.")
@dataclass

57
src/mflux/dreambooth/statistics/plotter.py
View File

@ -1,6 +1,3 @@
import time
from datetime import timedelta
import matplotlib.pyplot as plt
from mflux.dreambooth.state.training_spec import TrainingSpec
@ -8,8 +5,6 @@ from mflux.dreambooth.state.training_state import TrainingState
class Plotter:
start_time = time.time() # Class variable to track time
@staticmethod
def update_loss_plot(training_spec: TrainingSpec, training_state: TrainingState, target_loss: float = 0.3) -> None:
plt.style.use("bmh")
@ -22,32 +17,18 @@ class Plotter:
plt.plot(stats.steps, stats.losses, "b-", linewidth=2, label="Validation Loss", zorder=1) # Line
plt.plot(stats.steps, stats.losses, "bo", markersize=6, label="_nolegend_", zorder=2) # Points
# Find significant points - convert to Python float
# Find significant points
max_loss = float(max(stats.losses))
max_step = int(stats.steps[stats.losses.index(max_loss)])
min_loss = float(min(stats.losses))
min_step = int(stats.steps[stats.losses.index(min_loss)])
# Function to position annotations avoiding overlaps
def get_annotation_position(x, y, is_max=True, y_limits=None):
vertical_offset = 0.10
horizontal_offset = 0.10 # Horizontal offset to move to the left
if y_limits is None:
y_limits = plt.ylim()
if is_max:
if y + vertical_offset > y_limits[1]:
return float(x - horizontal_offset), float(y)
return float(x), float(y + vertical_offset)
else:
if y - vertical_offset < y_limits[0]:
return float(x - horizontal_offset), float(y)
return float(x), float(y - vertical_offset)
# Get the limits of the y-axis
y_limits = (0, 1.5) if max_loss <= 1.5 else (0, max_loss + max_loss * 0.2)
# Calculate y-axis limits with 20% padding above max loss
y_upper = max_loss + max_loss * 0.2
y_limits = (0, y_upper)
# Annotate for maximum and minimum loss
pos = get_annotation_position(max_step, max_loss, is_max=True, y_limits=y_limits)
pos = Plotter._get_annotation_position(max_step, max_loss, is_max=True, y_limits=y_limits)
plt.annotate(
f"Max: {max_loss:.3f}\nStep: {max_step}",
xy=(max_step, max_loss),
@ -59,7 +40,7 @@ class Plotter:
)
plt.plot(max_step, max_loss, "ro", markersize=8, zorder=4)
pos = get_annotation_position(min_step, min_loss, is_max=False, y_limits=y_limits)
pos = Plotter._get_annotation_position(min_step, min_loss, is_max=False, y_limits=y_limits)
plt.annotate(
f"Min: {min_loss:.3f}\nStep: {min_step}",
xy=(min_step, min_loss),
@ -74,10 +55,6 @@ class Plotter:
# Line for target loss
plt.axhline(y=float(target_loss), color="red", linestyle="--", linewidth=0.5, alpha=0.7)
# Calculate elapsed time
elapsed_time = time.time() - Plotter.start_time
elapsed_str = timedelta(seconds=int(elapsed_time))
# Initialize counter and sum of losses
loss_counter = 0
loss_sum = 0
@ -87,11 +64,9 @@ class Plotter:
loss_counter += 1
loss_sum += loss
if loss_counter > 0:
avg_loss = loss_sum / loss_counter
avg_loss = loss_sum / loss_counter if loss_counter > 0 else 0
legend_text = [
f"Elapsed Time: {elapsed_str}",
f"Img Dim {training_spec.width}x{training_spec.height}",
f"Total steps: {int(max(stats.steps))} / {total_step}",
f"Last Loss: {float(stats.losses[-1]):.4f}",
@ -122,11 +97,10 @@ class Plotter:
plt.subplots_adjust(right=0.85)
# Dinamiyc padding for x axes
# Dynamic padding for x axes
max_x = max(stats.steps)
initial_padding = 0.4
final_padding = 0.01
# formula for padding and centering graphos, it will progressivaly tent to final_padding.
padding_limit = initial_padding - (initial_padding - final_padding) * (max_x / total_step)
padding = max_x * padding_limit
@ -143,3 +117,18 @@ class Plotter:
plt.savefig(path, format="pdf", dpi=300, bbox_inches="tight")
plt.close()
@staticmethod
def _get_annotation_position(x, y, is_max=True, y_limits=None):
vertical_offset = 0.10
horizontal_offset = 0.10
if y_limits is None:
y_limits = plt.ylim()
if is_max:
if y + vertical_offset > y_limits[1]:
return float(x - horizontal_offset), float(y)
return float(x), float(y + vertical_offset)
else:
if y - vertical_offset < y_limits[0]:
return float(x - horizontal_offset), float(y)
return float(x), float(y - vertical_offset)