From 401da9fd75c1991b793a1650285aed4c187ceb2c Mon Sep 17 00:00:00 2001 From: Alessandro <73417515+azrahello@users.noreply.github.com> Date: Mon, 27 Jan 2025 23:12:17 +0100 Subject: [PATCH 1/2] Rebased on main (by filipstrand) --- .../dreambooth/lora_layers/lora_layers.py | 11 +- src/mflux/dreambooth/state/training_spec.py | 14 +- src/mflux/dreambooth/statistics/plotter.py | 120 ++++++++++++++++-- 3 files changed, 123 insertions(+), 22 deletions(-) diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py index 2923f77..ee6af18 100644 --- a/src/mflux/dreambooth/lora_layers/lora_layers.py +++ b/src/mflux/dreambooth/lora_layers/lora_layers.py @@ -70,12 +70,13 @@ class LoRALayers: blocks: list[JointTransformerBlock] | list[SingleTransformerBlock], block_prefix: str, ) -> dict: - start = block_spec.block_range.start - end = block_spec.block_range.end + block_indices = block_spec.block_range.get_blocks() lora_layers = {} - for i in range(start, end): - block = blocks[i] + for idx in block_indices: + if idx >= len(blocks): + raise IndexError(f"Indice {idx} over range") + block = blocks[idx] for layer_type in block_spec.layer_types: original_layer = LoRALayers._get_nested_attr(block, layer_type) is_list = isinstance(original_layer, list) @@ -84,7 +85,7 @@ class LoRALayers: linear=original_layer[0] if is_list else original_layer, r=block_spec.lora_rank, ) - layer_path = f"{block_prefix}.{i}.{layer_type}" + layer_path = f"{block_prefix}.{idx}.{layer_type}" lora_layers[layer_path] = [lora_layer] if is_list else lora_layer diff --git a/src/mflux/dreambooth/state/training_spec.py b/src/mflux/dreambooth/state/training_spec.py index 7438427..3ac38cc 100644 --- a/src/mflux/dreambooth/state/training_spec.py +++ b/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 +from typing import List, Optional from mflux.dreambooth.state.zip_util import ZipUtil @@ -56,8 +56,16 @@ class StatisticsSpec: @dataclass class BlockRange: - start: int - end: int + start: Optional[int] = None + end: Optional[int] = None + indices: Optional[List[int]] = 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.") @dataclass diff --git a/src/mflux/dreambooth/statistics/plotter.py b/src/mflux/dreambooth/statistics/plotter.py index 111d3c3..59d4ead 100644 --- a/src/mflux/dreambooth/statistics/plotter.py +++ b/src/mflux/dreambooth/statistics/plotter.py @@ -1,3 +1,6 @@ +import time +from datetime import timedelta + import matplotlib.pyplot as plt from mflux.dreambooth.state.training_spec import TrainingSpec @@ -5,8 +8,10 @@ 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) -> None: + def update_loss_plot(training_spec: TrainingSpec, training_state: TrainingState, target_loss: float = 0.3) -> None: plt.style.use("bmh") # Create figure with 16:9 aspect ratio @@ -17,37 +22,124 @@ 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 - # Customize the plot + # Find significant points - convert to Python float + 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) + + # Annotate for maximum and minimum loss + pos = 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), + xytext=pos, + bbox=dict(facecolor="red", alpha=0.5), + ha="center", + arrowprops=dict(arrowstyle="->"), + zorder=5, + ) + 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) + plt.annotate( + f"Min: {min_loss:.3f}\nStep: {min_step}", + xy=(min_step, min_loss), + xytext=pos, + bbox=dict(facecolor="green", alpha=0.5), + ha="center", + arrowprops=dict(arrowstyle="->"), + zorder=5, + ) + plt.plot(min_step, min_loss, "go", markersize=8, zorder=4) + + # 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 + total_step = int(training_state.iterator.total_number_of_steps()) + + for loss in stats.losses: + loss_counter += 1 + loss_sum += loss + + if loss_counter > 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}", + f"Lower Loss: {min_loss:.4f} (Step {min_step})", + f"Higher loss: {max_loss:.4f} (Step {max_step})", + f"Avg Loss: {avg_loss:.2f}", + f"Ideal Goal (lower of): {target_loss}", + ] + + plt.text( + 0.98, + 0.98, + "\n".join(legend_text), + transform=plt.gca().transAxes, + bbox=dict(facecolor="white", alpha=0.8, edgecolor="gray"), + verticalalignment="top", + horizontalalignment="right", + fontsize=8, + ) + plt.title("Validation Loss Over Time", fontsize=16, pad=20) plt.xlabel("Steps", fontsize=12) plt.ylabel("Loss", fontsize=12) - # Set integer grid plt.grid(True, linestyle="--", alpha=0.7) ax = plt.gca() ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True)) - # Dynamic x-axis limit with 20% padding + plt.subplots_adjust(right=0.85) + + # Dinamiyc padding for x axes max_x = max(stats.steps) - padding = max_x * 0.2 + 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 + plt.xlim(0, max_x + padding) - # Dynamic y-axis limit with 20% padding - max_y = float(max(stats.losses)) - padding = max_y * 0.2 - plt.ylim(0, max_y + padding) - + plt.ylim(y_limits) plt.legend(fontsize=12) - # Add margins for better visibility plt.margins(x=0.02) - # Tight layout to prevent label cutoff plt.tight_layout() - # Save to desktop with high PPI path = training_state.get_current_loss_plot_path(training_spec) plt.savefig(path, format="pdf", dpi=300, bbox_inches="tight") - # Close the figure to free memory plt.close() From f94edc4c3b2a0621d9cb102fe5f61e8484c8fb62 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Mon, 27 Jan 2025 22:39:29 +0100 Subject: [PATCH 2/2] Some fixes: Removed time in plotting, fix off-by-one error in training spec, update y-axis on loss plot --- README.md | 24 +++++++- .../dreambooth/lora_layers/lora_layers.py | 2 +- src/mflux/dreambooth/state/training_spec.py | 12 ++-- src/mflux/dreambooth/statistics/plotter.py | 57 ++++++++----------- 4 files changed, 53 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index b204643..b2c85dd 100644 --- a/README.md +++ b/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` +
+Specify individual layers + +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 + ], + ... + }, +... +``` +
+ *⚠️ 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.* diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py index ee6af18..2ad62d4 100644 --- a/src/mflux/dreambooth/lora_layers/lora_layers.py +++ b/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: diff --git a/src/mflux/dreambooth/state/training_spec.py b/src/mflux/dreambooth/state/training_spec.py index 3ac38cc..a98e996 100644 --- a/src/mflux/dreambooth/state/training_spec.py +++ b/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 - 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 diff --git a/src/mflux/dreambooth/statistics/plotter.py b/src/mflux/dreambooth/statistics/plotter.py index 59d4ead..1bb4929 100644 --- a/src/mflux/dreambooth/statistics/plotter.py +++ b/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 - 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)