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 2923f77..2ad62d4 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"Index {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..a98e996 100644 --- a/src/mflux/dreambooth/state/training_spec.py +++ b/src/mflux/dreambooth/state/training_spec.py @@ -56,8 +56,16 @@ class StatisticsSpec: @dataclass class BlockRange: - start: int - end: int + 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)) + 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..1bb4929 100644 --- a/src/mflux/dreambooth/statistics/plotter.py +++ b/src/mflux/dreambooth/statistics/plotter.py @@ -6,7 +6,7 @@ from mflux.dreambooth.state.training_state import TrainingState class Plotter: @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 +17,118 @@ 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 + 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)]) + + # 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 = 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), + 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 = 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), + 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) + + # 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 + + avg_loss = loss_sum / loss_counter if loss_counter > 0 else 0 + + legend_text = [ + 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) + + # Dynamic padding for x axes max_x = max(stats.steps) - padding = max_x * 0.2 + initial_padding = 0.4 + final_padding = 0.01 + 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() + + @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)