autoforge 1.8.0__tar.gz → 1.8.2__tar.gz

{autoforge-1.8.0 → autoforge-1.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.8.0
+Version: 1.8.2
 Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
 Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
 License: CC BY-NC-SA 4.0
@@ -66,22 +66,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
 - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
 - **Hueforge export**: Outputs a project file that can be opened with hueforge.
 
-## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
-
-If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
-
-1. **Open the Notebook:**  
-   Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
-
-2. **Select GPU Runtime:**  
-   In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
-
-3. **Run All Cells:**  
-   Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
-
-4. **Review and Download:**  
-   After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
-
 ## Manual Installation
 
 To install AutoForge, simply install the current version from PyPI:
@@ -116,7 +100,6 @@ autoforge --input_image path/to/input_image.jpg --json_file %APPDATA%\HueForge\F
 autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
 ```
 
-
 ### Command Line Arguments
 
 - `--config` *(Optional)* Path to a configuration file with the settings.

{autoforge-1.8.0 → autoforge-1.8.2}/README.md

@@ -40,22 +40,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
 - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
 - **Hueforge export**: Outputs a project file that can be opened with hueforge.
 
-## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
-
-If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
-
-1. **Open the Notebook:**  
-   Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
-
-2. **Select GPU Runtime:**  
-   In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
-
-3. **Run All Cells:**  
-   Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
-
-4. **Review and Download:**  
-   After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
-
 ## Manual Installation
 
 To install AutoForge, simply install the current version from PyPI:
@@ -90,7 +74,6 @@ autoforge --input_image path/to/input_image.jpg --json_file %APPDATA%\HueForge\F
 autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
 ```
 
-
 ### Command Line Arguments
 
 - `--config` *(Optional)* Path to a configuration file with the settings.

{autoforge-1.8.0 → autoforge-1.8.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "AutoForge"
-version="1.8.0"
+version="1.8.2"
 description = "AutoForge is a Python tool for generating 3D printed layered models from an input image."
 authors = [{ name = "Hendric Voss", email = "hvoss@techfak.uni-bielefeld.de" }]
 readme = "README.md"

{autoforge-1.8.0 → autoforge-1.8.2}/src/AutoForge.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.8.0
+Version: 1.8.2
 Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
 Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
 License: CC BY-NC-SA 4.0
@@ -66,22 +66,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
 - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
 - **Hueforge export**: Outputs a project file that can be opened with hueforge.
 
-## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
-
-If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
-
-1. **Open the Notebook:**  
-   Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
-
-2. **Select GPU Runtime:**  
-   In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
-
-3. **Run All Cells:**  
-   Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
-
-4. **Review and Download:**  
-   After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
-
 ## Manual Installation
 
 To install AutoForge, simply install the current version from PyPI:
@@ -116,7 +100,6 @@ autoforge --input_image path/to/input_image.jpg --json_file %APPDATA%\HueForge\F
 autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
 ```
 
-
 ### Command Line Arguments
 
 - `--config` *(Optional)* Path to a configuration file with the settings.

{autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/OptimizerHelper.py

@@ -277,38 +277,3 @@ def composite_image_disc(
     comp = comp + rem_after.unsqueeze(-1) * background  # [H,W,3]
 
     return comp * 255.0
-
-
-def discretize_solution(
-    params: dict, tau_global: float, h: float, max_layers: int, rng_seed: int = -1
-):
-    """
-    Convert continuous logs to discrete layer counts and discrete color IDs.
-
-    Args:
-        params (dict): Dictionary containing the parameters 'pixel_height_logits' and 'global_logits'.
-        tau_global (float): Temperature parameter for global material assignment.
-        h (float): Height of each layer.
-        max_layers (int): Maximum number of layers.
-        rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
-
-    Returns:
-        tuple: A tuple containing:
-            - torch.Tensor: Discrete global material assignments, shape [max_layers].
-            - torch.Tensor: Discrete height image, shape [H, W].
-    """
-    pixel_height_logits = params["pixel_height_logits"]
-    global_logits = params["global_logits"]
-    pixel_heights = (max_layers * h) * torch.sigmoid(pixel_height_logits)
-    discrete_height_image = torch.round(pixel_heights / h).to(torch.int32)
-    discrete_height_image = torch.clamp(discrete_height_image, 0, max_layers)
-
-    num_layers = global_logits.shape[0]
-    discrete_global_vals = []
-    for j in range(num_layers):
-        p = deterministic_gumbel_softmax(
-            global_logits[j], tau_global, hard=True, rng_seed=rng_seed + j
-        )
-        discrete_global_vals.append(torch.argmax(p))
-    discrete_global = torch.stack(discrete_global_vals, dim=0)
-    return discrete_global, discrete_height_image

{autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/PruningHelper.py

@@ -6,7 +6,7 @@ from tqdm import tqdm
 from joblib import Parallel, delayed
 import threading
 
-from autoforge.Helper.OptimizerHelper import discretize_solution, composite_image_disc
+from autoforge.Helper.OptimizerHelper import composite_image_disc
 from autoforge.Loss.LossFunctions import compute_loss
 from autoforge.Modules.Optimizer import FilamentOptimizer
 
@@ -325,7 +325,13 @@ def merge_bands(
 
 
 def remove_layer_from_solution(
-    params, layer_to_remove, final_tau, h, current_max_layers, rng_seed
+    optimizer,
+    params,
+    layer_to_remove,
+    final_tau,
+    h,
+    current_max_layers,
+    rng_seed,
 ):
     """
     Remove one layer from the solution.
@@ -343,8 +349,15 @@ def remove_layer_from_solution(
         new_max_layers (int): Updated number of layers.
     """
     # Get the current discrete height image (used to decide which pixels need adjusting)
-    _, disc_height = discretize_solution(
-        params, final_tau, h, current_max_layers, rng_seed
+    effective_logits = optimizer._apply_height_offset(
+        params["pixel_height_logits"], params["height_offsets"]
+    )
+    _, disc_height = optimizer.discretize_solution(
+        params,
+        final_tau,
+        h,
+        current_max_layers,
+        rng_seed,
     )
 
     # Remove the candidate layer from the global (color) assignment.
@@ -358,25 +371,24 @@ def remove_layer_from_solution(
     new_max_layers = new_global_logits.shape[0]
 
     # Compute current effective height: height = (current_max_layers * h) * sigmoid(pixel_height_logits)
-    current_height = (
-        current_max_layers * h * torch.sigmoid(params["pixel_height_logits"])
-    )
-
-    # For pixels where the discrete height is at or above the removed layer, subtract h.
+    current_height = current_max_layers * h * torch.sigmoid(effective_logits)
     new_height = current_height.clone()
-    mask = disc_height >= layer_to_remove  # <-- Changed from > to >= here
+    mask = disc_height >= layer_to_remove
     new_height[mask] = new_height[mask] - h
 
     # Invert the sigmoid mapping:
     # We need new_pixel_height_logits such that:
     #    sigmoid(new_pixel_height_logits) = new_height / (new_max_layers * h)
     eps = 1e-6
-    new_ratio = torch.clamp(new_height / (new_max_layers * h), eps, 1 - eps)
-    new_pixel_height_logits = torch.log(new_ratio) - torch.log(1 - new_ratio)
-
+    ratio = torch.clamp(new_height / (new_max_layers * h), eps, 1.0 - eps)
+    new_effective_logits = torch.log(ratio) - torch.log1p(-ratio)
+    new_pixel_height_logits = new_effective_logits - (
+        effective_logits - params["pixel_height_logits"]
+    )
     new_params = {
         "global_logits": new_global_logits,
         "pixel_height_logits": new_pixel_height_logits,
+        "height_offsets": params["height_offsets"],
     }
     return new_params, new_max_layers
 
@@ -399,129 +411,157 @@ def prune_redundant_layers(
     stopping as soon as one batch yields an improvement.
     """
     current_max_layers = optimizer.best_params["global_logits"].shape[0]
-    optimizer.max_layers = current_max_layers
-
-    def score_layer(
-        layer_idx: int,
-        base_params: dict,
-        final_tau: float,
-        h: float,
-        cur_max_layers: int,
-        best_seed: int,
-    ):
+    optimizer.max_layers = current_max_layers  # keep optimiser in sync
+
+    # Baseline loss with current best parameters
+    with _gpu_lock, torch.no_grad():
+        eff_logits = optimizer._apply_height_offset(
+            optimizer.best_params["pixel_height_logits"],
+            optimizer.best_params["height_offsets"],
+        )
+        ref_comp = composite_image_disc(
+            eff_logits,
+            optimizer.best_params["global_logits"],
+            optimizer.vis_tau,
+            optimizer.vis_tau,
+            optimizer.h,
+            current_max_layers,
+            optimizer.material_colors,
+            optimizer.material_TDs,
+            optimizer.background,
+            rng_seed=optimizer.best_seed,
+        )
+        best_loss = compute_loss(ref_comp, optimizer.target).item()
+
+    tbar = tqdm(
+        desc=f"Layer pruning | Loss {best_loss:.4f}",
+        total=max(current_max_layers - 1, 0),
+        leave=False,
+    )
+    removed_layers = 0
+    improvement = True
+
+    # ----------------------------------------------------------
+    # Inner helper: evaluate the effect of removing layer idx
+    # ----------------------------------------------------------
+    def score_layer(idx: int) -> tuple[float, dict, int]:
         cand_params, cand_max_layers = remove_layer_from_solution(
-            base_params, layer_idx, final_tau, h, cur_max_layers, best_seed
+            optimizer,
+            optimizer.best_params,
+            idx,
+            optimizer.vis_tau,
+            optimizer.h,
+            current_max_layers,
+            optimizer.best_seed,
+        )
+        eff_logits = optimizer._apply_height_offset(
+            cand_params["pixel_height_logits"], cand_params["height_offsets"]
         )
         with _gpu_lock, torch.no_grad():
             cand_comp = composite_image_disc(
-                cand_params["pixel_height_logits"],
+                eff_logits,
                 cand_params["global_logits"],
-                final_tau,
-                final_tau,
-                h,
+                optimizer.vis_tau,
+                optimizer.vis_tau,
+                optimizer.h,
                 cand_max_layers,
                 optimizer.material_colors,
                 optimizer.material_TDs,
                 optimizer.background,
-                rng_seed=best_seed,
+                rng_seed=optimizer.best_seed,
             )
             cand_loss = compute_loss(cand_comp, optimizer.target).item()
         return cand_loss, cand_params, cand_max_layers
 
-    with _gpu_lock, torch.no_grad():
-        comp = optimizer.get_best_discretized_image()
-        best_loss = compute_loss(comp, optimizer.target).item()
-
-    tbar = tqdm(
-        desc=f"Layer pruning | Loss {best_loss:.4f}",
-        total=max(current_max_layers - 1, 0),
-        leave=False,
-    )
-    removed_layers = 0
-    improvement = True
-
+    # ----------------------------------------------------------
+    # Main pruning loop
+    # ----------------------------------------------------------
     while current_max_layers > pruning_min_layers and (
         improvement or current_max_layers > pruning_max_layers
     ):
         tbar.update(1)
         improvement = False
-
         layer_indices = list(range(current_max_layers))
+
         if fast:
             random.shuffle(layer_indices)
             chunk_size = max(1, math.ceil(len(layer_indices) * chunking_percent))
-            c_cand = None
-            c_loss = 1000
+            best_candidate = None
+            best_cand_loss = float("inf")
+
             for chunk in _chunked(layer_indices, chunk_size):
                 cand_results = Parallel(
                     n_jobs=n_jobs, backend="threading", prefer="threads"
-                )(
-                    delayed(score_layer)(
-                        idx,
-                        optimizer.best_params,
-                        optimizer.vis_tau,
-                        optimizer.h,
-                        current_max_layers,
-                        optimizer.best_seed,
-                    )
-                    for idx in chunk
-                )
+                )(delayed(score_layer)(idx) for idx in chunk)
+
                 cand_loss, cand_params, cand_max_layers = min(
                     cand_results, key=lambda x: x[0]
                 )
+
                 if cand_loss <= best_loss * (1 + allowed_loss_increase_percent):
+                    # Accept immediately and restart outer loop
                     removed_layers += 1
                     best_loss = cand_loss
                     current_max_layers = cand_max_layers
                     optimizer.best_params = cand_params
                     optimizer.max_layers = current_max_layers
                     tbar.set_description(
-                        f"Layer pruning | Loss {best_loss:.4f} | Removed {removed_layers} | Layers {current_max_layers}"
+                        f"Layer pruning | Loss {best_loss:.4f} | "
+                        f"Removed {removed_layers} | Layers {current_max_layers}"
                     )
                     improvement = True
-                    break  # restart outer while loop with updated solution
-                else:
-                    if current_max_layers > pruning_max_layers and cand_loss < c_loss:
-                        c_cand = (cand_params, current_max_layers)
-            if c_cand is not None:
+                    break
+
+                # Track the best candidate in case we must enforce a hard limit
+                if cand_loss < best_cand_loss:
+                    best_cand_loss = cand_loss
+                    best_candidate = (cand_params, cand_max_layers)
+
+            # Forced removal to meet pruning_max_layers
+            if (
+                not improvement
+                and current_max_layers > pruning_max_layers
+                and best_candidate is not None
+            ):
+                cand_params, cand_max_layers = best_candidate
                 removed_layers += 1
-                best_loss = c_loss
-                current_max_layers = c_cand[1]
-                optimizer.best_params = c_cand[0]
+                best_loss = best_cand_loss
+                current_max_layers = cand_max_layers
+                optimizer.best_params = cand_params
                 optimizer.max_layers = current_max_layers
                 tbar.set_description(
-                    f"Layer pruning | Loss {best_loss:.4f} | Removed {removed_layers} | Layers {current_max_layers}"
+                    f"Layer pruning | Loss {best_loss:.4f} | "
+                    f"Removed {removed_layers} | Layers {current_max_layers}"
                 )
                 improvement = True
 
             if not improvement:
-                break
+                break  # no chunk helped and layer count is within budget
         else:
+            # Exhaustive search
             cand_results = Parallel(
                 n_jobs=n_jobs, backend="threading", prefer="threads"
-            )(
-                delayed(score_layer)(
-                    idx,
-                    optimizer.best_params,
-                    optimizer.vis_tau,
-                    optimizer.h,
-                    current_max_layers,
-                    optimizer.best_seed,
-                )
-                for idx in layer_indices
-            )
+            )(delayed(score_layer)(idx) for idx in layer_indices)
+
             cand_loss, cand_params, cand_max_layers = min(
                 cand_results, key=lambda x: x[0]
             )
-            if cand_loss <= best_loss or current_max_layers > pruning_max_layers:
+
+            tbar.set_description(
+                f"Layer pruning | Current Candidate Loss: {cand_loss:.4f} | Best Loss {best_loss:.4f} | "
+                f"Removed {removed_layers} | Layers {current_max_layers}"
+            )
+
+            if (
+                cand_loss <= best_loss * (1 + allowed_loss_increase_percent)
+                or current_max_layers > pruning_max_layers
+            ):
                 removed_layers += 1
                 best_loss = cand_loss
                 current_max_layers = cand_max_layers
                 optimizer.best_params = cand_params
                 optimizer.max_layers = current_max_layers
-                tbar.set_description(
-                    f"Layer pruning | Loss {best_loss:.4f} | Removed {removed_layers} | Layers {current_max_layers}"
-                )
+
                 improvement = True
             else:
                 break

{autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Modules/Optimizer.py

@@ -9,9 +9,9 @@ from tqdm import tqdm
 
 from autoforge.Helper.CAdamW import CAdamW
 from autoforge.Helper.OptimizerHelper import (
-    discretize_solution,
     composite_image_cont,
     composite_image_disc,
+    deterministic_gumbel_softmax,
 )
 from autoforge.Loss.LossFunctions import loss_fn, compute_loss
 
@@ -115,6 +115,7 @@ class FilamentOptimizer:
         self.params = {
             "pixel_height_logits": self.pixel_height_logits,
             "global_logits": global_logits_init,
+            "height_offsets": self.height_offsets,
         }
 
         # Tau schedule
@@ -303,6 +304,50 @@ class FilamentOptimizer:
 
         return loss
 
+    def discretize_solution(
+        self,
+        params: dict,
+        tau_global: float,
+        h: float,
+        max_layers: int,
+        rng_seed: int = -1,
+    ):
+        """
+        Convert continuous logs to discrete layer counts and discrete color IDs.
+
+        Args:
+            params (dict): Dictionary containing the parameters 'pixel_height_logits' and 'global_logits'.
+            tau_global (float): Temperature parameter for global material assignment.
+            h (float): Height of each layer.
+            max_layers (int): Maximum number of layers.
+            rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
+
+        Returns:
+            tuple: A tuple containing:
+                - torch.Tensor: Discrete global material assignments, shape [max_layers].
+                - torch.Tensor: Discrete height image, shape [H, W].
+        """
+        pixel_logits = params["pixel_height_logits"]
+        pixel_offset = params["height_offsets"]
+        effective_logits = self._apply_height_offset(
+            pixel_logits=pixel_logits, height_offsets=pixel_offset
+        )
+
+        global_logits = params["global_logits"]
+        pixel_heights = (max_layers * h) * torch.sigmoid(effective_logits)
+        discrete_height_image = torch.round(pixel_heights / h).to(torch.int32)
+        discrete_height_image = torch.clamp(discrete_height_image, 0, max_layers)
+
+        num_layers = global_logits.shape[0]
+        discrete_global_vals = []
+        for j in range(num_layers):
+            p = deterministic_gumbel_softmax(
+                global_logits[j], tau_global, hard=True, rng_seed=rng_seed + j
+            )
+            discrete_global_vals.append(torch.argmax(p))
+        discrete_global = torch.stack(discrete_global_vals, dim=0)
+        return discrete_global, discrete_height_image
+
     def log_to_tensorboard(
         self, interval: int = 100, namespace: str = "", step: int = None
     ):
@@ -489,13 +534,8 @@ class FilamentOptimizer:
             )
 
         if best:
-            effective_logits = self._apply_height_offset(
-                self.best_params["pixel_height_logits"],
-                self.best_params["height_offsets"],
-            )
-            current_params["pixel_height_logits"] = effective_logits
-            disc_global, disc_height_image = discretize_solution(
-                current_params,
+            disc_global, disc_height_image = self.discretize_solution(
+                self.best_params,
                 self.vis_tau,
                 self.h,
                 self.max_layers,
@@ -505,7 +545,7 @@ class FilamentOptimizer:
         else:
             tau_height, tau_global = self._get_tau()
             with torch.no_grad():
-                disc_global, disc_height_image = discretize_solution(
+                disc_global, disc_height_image = self.discretize_solution(
                     current_params,
                     tau_height,
                     self.h,
@@ -603,18 +643,14 @@ class FilamentOptimizer:
             tau_h, tau_g = self.vis_tau, self.vis_tau
             with torch.no_grad():
                 effective_logits = self._apply_height_offset()
-                params_with_offset = {
-                    "pixel_height_logits": effective_logits,
-                    "global_logits": self.params["global_logits"],
-                }
-                disc_global, disc_height_image = discretize_solution(
-                    params_with_offset, tau_g, self.h, self.max_layers, rng_seed=seed
+                disc_global, disc_height_image = self.discretize_solution(
+                    self.params, tau_g, self.h, self.max_layers, rng_seed=seed
                 )
 
                 # 2) Compute discrete-mode composite
                 with torch.no_grad():
                     comp_disc = composite_image_disc(
-                        self.params["pixel_height_logits"],
+                        effective_logits,
                         self.params["global_logits"],
                         self.vis_tau,
                         self.vis_tau,
@@ -659,8 +695,12 @@ class FilamentOptimizer:
         best_loss = start_loss
         for i in tqdm(range(num_seeds), desc="Searching for new best seed"):
             seed = np.random.randint(0, 1000000)
-            comp_disc = composite_image_disc(
+            effective_logits = self._apply_height_offset(
                 self.best_params["pixel_height_logits"],
+                self.best_params["height_offsets"],
+            )
+            comp_disc = composite_image_disc(
+                effective_logits,
                 self.best_params["global_logits"],
                 self.vis_tau,
                 self.vis_tau,