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

README.md

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

src/mflux/dreambooth/lora_layers/lora_layers.py

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

src/mflux/dreambooth/state/training_spec.py

@@ -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
+    start: int | None = None
-    end: Optional[int] = None
+    end: int | None = None
-    indices: Optional[List[int]] = 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))
+            return list(range(self.start, self.end))
-        raise ValueError("Either ‘start’ and ‘end’ or ‘indices’ must be provided.")
+        raise ValueError("Either 'start' and 'end' or 'indices' must be provided.")
 
 
 @dataclass

src/mflux/dreambooth/statistics/plotter.py

@@ -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
+        # Calculate y-axis limits with 20% padding above max loss
-        def get_annotation_position(x, y, is_max=True, y_limits=None):
+        y_upper = max_loss + max_loss * 0.2
-            vertical_offset = 0.10
+        y_limits = (0, y_upper)
-            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)
 
         # 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 if loss_counter > 0 else 0
-            avg_loss = loss_sum / loss_counter
 
         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)