autoforge 1.4.2__tar.gz → 1.5.0__tar.gz

{autoforge-1.4.2 → autoforge-1.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.4.2
+Version: 1.5.0
 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

{autoforge-1.4.2 → autoforge-1.5.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "AutoForge"
-version="1.4.2"
+version="1.5.0"
 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.4.2 → autoforge-1.5.0}/src/AutoForge.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.4.2
+Version: 1.5.0
 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

{autoforge-1.4.2 → autoforge-1.5.0}/src/AutoForge.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ src/autoforge/Gui/main.py
 src/autoforge/Gui/Views/__init__.py
 src/autoforge/Gui/Views/backend_database.py
 src/autoforge/Gui/Views/filament_editor.py
+src/autoforge/Helper/CAdamW.py
 src/autoforge/Helper/FilamentHelper.py
 src/autoforge/Helper/ImageHelper.py
 src/autoforge/Helper/OptimizerHelper.py

autoforge-1.5.0/src/autoforge/Helper/CAdamW.py

@@ -0,0 +1,113 @@
+# Taken from https://github.com/kyleliang919/C-Optim/blob/main/c_adamw.py
+# copy dependencies from transformers/optimization.py
+import math
+import warnings
+from typing import Callable, Iterable, Tuple
+
+import torch
+from torch import nn
+from torch.optim import Optimizer
+
+
+class CAdamW(Optimizer):
+    """
+    Implements Adam algorithm with weight decay fix as introduced in [Decoupled Weight Decay
+    Regularization](https://arxiv.org/abs/1711.05101).
+
+    Parameters:
+        params (`Iterable[nn.parameter.Parameter]`):
+            Iterable of parameters to optimize or dictionaries defining parameter groups.
+        lr (`float`, *optional*, defaults to 0.001):
+            The learning rate to use.
+        betas (`Tuple[float,float]`, *optional*, defaults to `(0.9, 0.999)`):
+            Adam's betas parameters (b1, b2).
+        eps (`float`, *optional*, defaults to 1e-06):
+            Adam's epsilon for numerical stability.
+        weight_decay (`float`, *optional*, defaults to 0.0):
+            Decoupled weight decay to apply.
+        correct_bias (`bool`, *optional*, defaults to `True`):
+            Whether or not to correct bias in Adam (for instance, in Bert TF repository they use `False`).
+        no_deprecation_warning (`bool`, *optional*, defaults to `False`):
+            A flag used to disable the deprecation warning (set to `True` to disable the warning).
+    """
+
+    def __init__(
+            self,
+            params: Iterable[nn.parameter.Parameter],
+            lr: float = 1e-3,
+            betas: Tuple[float, float] = (0.9, 0.999),
+            eps: float = 1e-6,
+            weight_decay: float = 0.0,
+            correct_bias: bool = True,
+    ):
+
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps} - should be >= 0.0")
+        defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "correct_bias": correct_bias}
+        super().__init__(params, defaults)
+        self.init_lr = lr
+
+    @torch.no_grad()
+    def step(self, closure: Callable = None):
+        """
+        Performs a single optimization step.
+
+        Arguments:
+            closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for i, p in enumerate(group["params"]):
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                state = self.state[p]
+
+                if "step" not in state:
+                    state["step"] = 0
+
+                # State initialization
+                if "exp_avg" not in state:
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(grad)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(grad)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                # apply weight decay
+                if group["weight_decay"] > 0.0:
+                    p.add_(p, alpha=(-group["lr"] * group["weight_decay"]))
+
+                # Decay the first and second moment running average coefficient
+                # In-place operations to update the averages at the same time
+                exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+                denom = exp_avg_sq.sqrt().add_(group["eps"])
+
+                step_size = group["lr"]
+                if group["correct_bias"]:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1 ** state["step"]
+                    bias_correction2 = 1.0 - beta2 ** state["step"]
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                # compute norm gradient
+                mask = (exp_avg * grad > 0).to(grad.dtype)
+                # mask = mask * (mask.numel() / (mask.sum() + 1)) ## original implementation, leaving it here for record
+                mask.div_(mask.mean().clamp_(min=1e-3))  # https://huggingface.co/rwightman/timm-optim-caution found this implementation is more favoarable in many cases
+                norm_grad = (exp_avg * mask) / denom
+                p.add_(norm_grad, alpha=-step_size)
+        return loss

{autoforge-1.4.2 → autoforge-1.5.0}/src/autoforge/Helper/OptimizerHelper.py

@@ -76,7 +76,7 @@ def deterministic_gumbel_softmax(
 
 
 @torch.jit.script
-def composite_image(
+def composite_image_cont(
     pixel_height_logits: torch.Tensor,
     global_logits: torch.Tensor,
     tau_height: float,
@@ -86,32 +86,27 @@ def composite_image(
     material_colors: torch.Tensor,
     material_TDs: torch.Tensor,
     background: torch.Tensor,
-    mode: str = "continuous",
-    rng_seed: int = -1,  # <--- new optional argument
 ) -> torch.Tensor:
     """
-    Vectorized compositing over all pixels (H x W).
-    ...
-    If rng_seed >= 0 and mode=="discrete", we will fix the random seed
-    for each layer so that we get the same random Gumbel noise for reproducibility.
+    Continuous compositing over all pixels.
+    Uses Gumbel softmax with either hard or soft sampling depending on tau_global.
     """
     pixel_height = (max_layers * h) * torch.sigmoid(pixel_height_logits)
     continuous_layers = pixel_height / h
 
-    # same as before
     adaptive_layers = adaptive_round(
         continuous_layers, tau_height, high_tau=0.1, low_tau=0.01, temp=0.1
     )
     discrete_layers_temp = torch.round(continuous_layers)
     discrete_layers = (
         discrete_layers_temp + (adaptive_layers - discrete_layers_temp).detach()
-    )
-    discrete_layers = discrete_layers.to(torch.int32)
+    ).to(torch.int32)
 
     H, W = pixel_height.shape
     comp = torch.zeros(H, W, 3, dtype=torch.float32, device=pixel_height.device)
     remaining = torch.ones(H, W, dtype=torch.float32, device=pixel_height.device)
 
+    # opacity function parameters
     o = -1.2416557e-02
     A = 9.6407950e-01
     k = 3.4103447e01
@@ -119,39 +114,20 @@ def composite_image(
 
     for i in range(max_layers):
         layer_idx = max_layers - 1 - i
-        p_print = (discrete_layers > layer_idx).float()
-        eff_thick = p_print * h
-
-        # This is the key difference for "discrete":
-        # we now optionally fix the random seed if rng_seed >= 0
-        if mode == "discrete":
-            if rng_seed >= 0:
-                # Use a context manager so we don’t pollute the global seed
-                p_i = deterministic_gumbel_softmax(
-                    global_logits[layer_idx],
-                    tau_global,
-                    hard=True,
-                    rng_seed=rng_seed + layer_idx,
-                )
-            else:
-                # same as before, but will be random every call
-                p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=True)
-
-        elif mode == "continuous":
-            if tau_global < 1e-3:
-                p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=True)
-            else:
-                p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=False)
-        else:
-            # e.g. "pruning" or any other fallback
+        p_print = (discrete_layers > layer_idx).float()  # [H,W]
+        eff_thick = p_print * h  # effective thickness per pixel
+
+        # continuous mode: use hard gumbel if tau_global is very low
+        if tau_global < 1e-3:
             p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=True)
+        else:
+            p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=False)
 
         color_i = torch.matmul(p_i, material_colors)
         TD_i = torch.matmul(p_i, material_TDs)
         TD_i = torch.clamp(TD_i, 1e-8, 1e8)
 
         opac = o + (A * torch.log1p(k * (eff_thick / TD_i)) + b * (eff_thick / TD_i))
-        # opac = torch.where(TD_i / 10 <= h, torch.tensor(1.0), opac)
         opac = torch.clamp(opac, 0.0, 1.0)
 
         comp = comp + ((remaining * opac).unsqueeze(-1) * color_i)
@@ -161,6 +137,143 @@ def composite_image(
     return comp * 255.0
 
 
+@torch.jit.script
+def composite_image_disc(
+    pixel_height_logits: torch.Tensor,
+    global_logits: torch.Tensor,
+    tau_height: float,
+    tau_global: float,
+    h: float,
+    max_layers: int,
+    material_colors: torch.Tensor,
+    material_TDs: torch.Tensor,
+    background: torch.Tensor,
+    rng_seed: int = -1,
+) -> torch.Tensor:
+    """
+    Discrete compositing over all pixels.
+    For each layer (traversed from top to bottom), we compute a one-hot material selection.
+    For pixels where the layer is printed (as determined by discrete_layers),
+    contiguous layers with the same material are merged (i.e. their effective thickness is accumulated)
+    and then composited using the opacity function. The accumulated segment is applied when
+    the material changes or if the layer is not printed.
+    """
+    # Compute pixel height and determine discrete layers.
+    pixel_height = (max_layers * h) * torch.sigmoid(pixel_height_logits)
+    continuous_layers = pixel_height / h
+
+    adaptive_layers = adaptive_round(
+        continuous_layers, tau_height, high_tau=0.1, low_tau=0.01, temp=0.1
+    )
+    discrete_layers_temp = torch.round(continuous_layers)
+    discrete_layers = (
+        discrete_layers_temp + (adaptive_layers - discrete_layers_temp).detach()
+    ).to(torch.int32)
+
+    H, W = pixel_height.shape
+    comp = torch.zeros(H, W, 3, dtype=torch.float32, device=pixel_height.device)
+    remaining = torch.ones(H, W, dtype=torch.float32, device=pixel_height.device)
+
+    # Current material and accumulated thickness per pixel.
+    # cur_mat is initialized to -1 meaning "no material".
+    cur_mat = -torch.ones((H, W), dtype=torch.int32, device=pixel_height.device)
+    acc_thick = torch.zeros((H, W), dtype=torch.float32, device=pixel_height.device)
+
+    # Opacity function parameters.
+    o = -1.2416557e-02
+    A = 9.6407950e-01
+    k = 3.4103447e01
+    b = -4.1554203e00
+
+    for i in range(max_layers):
+        layer_idx = max_layers - 1 - i
+        p_print = discrete_layers > layer_idx  # Boolean mask [H,W]
+        eff_thick = (
+            p_print.to(torch.float32) * h
+        )  # effective thickness for printed pixels
+
+        # Determine material for this layer using discrete (one-hot) gumbel softmax.
+        if rng_seed >= 0:
+            p_i = deterministic_gumbel_softmax(
+                global_logits[layer_idx], tau_global, True, rng_seed + layer_idx
+            )
+        else:
+            p_i = F.gumbel_softmax(global_logits[layer_idx], tau_global, hard=True)
+        # Cast new_mat to int32 so it matches cur_mat.
+        new_mat = torch.argmax(p_i, 0).to(torch.int32)
+
+        # 1. For pixels where no current material is set, start a new segment.
+        mask_new = p_print & (cur_mat == -1)
+        if mask_new.any():
+            cur_mat[mask_new] = new_mat
+            acc_thick[mask_new] = eff_thick[mask_new]
+
+        # 2. For pixels where the current segment has the same material, accumulate thickness.
+        mask_same = p_print & (cur_mat == new_mat)
+        if mask_same.any():
+            acc_thick[mask_same] = acc_thick[mask_same] + eff_thick[mask_same]
+
+        # 3. For pixels where the current segment has a different material,
+        # composite the accumulated segment and start a new one.
+        mask_diff = p_print & (cur_mat != -1) & (cur_mat != new_mat)
+        if mask_diff.any():
+            # Convert indices to int64 for index_select.
+            indices_diff = cur_mat[mask_diff].to(torch.int64)
+            td_vals = material_TDs.index_select(0, indices_diff)
+            col_vals = material_colors.index_select(0, indices_diff)
+            opac_vals = o + (
+                A * torch.log1p(k * (acc_thick[mask_diff] / td_vals))
+                + b * (acc_thick[mask_diff] / td_vals)
+            )
+            opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
+            comp[mask_diff] = comp[mask_diff] + (
+                (remaining[mask_diff] * opac_vals).unsqueeze(-1) * col_vals
+            )
+            remaining[mask_diff] = remaining[mask_diff] * (1 - opac_vals)
+            cur_mat[mask_diff] = new_mat
+            acc_thick[mask_diff] = eff_thick[mask_diff]
+
+        # 4. For pixels where this layer is not printed but we have a pending segment,
+        # composite that segment and reset.
+        mask_not_printed = (~p_print) & (cur_mat != -1)
+        if mask_not_printed.any():
+            indices_np = cur_mat[mask_not_printed].to(torch.int64)
+            td_vals = material_TDs.index_select(0, indices_np)
+            col_vals = material_colors.index_select(0, indices_np)
+            opac_vals = o + (
+                A * torch.log1p(k * (acc_thick[mask_not_printed] / td_vals))
+                + b * (acc_thick[mask_not_printed] / td_vals)
+            )
+            opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
+            comp[mask_not_printed] = comp[mask_not_printed] + (
+                (remaining[mask_not_printed] * opac_vals).unsqueeze(-1) * col_vals
+            )
+            remaining[mask_not_printed] = remaining[mask_not_printed] * (1 - opac_vals)
+            cur_mat[mask_not_printed] = -1
+            acc_thick[mask_not_printed] = 0.0
+
+    # After the loop, composite any remaining accumulated segment.
+    mask_remain = cur_mat != -1
+    if mask_remain.any():
+        indices_rem = cur_mat[mask_remain].to(torch.int64)
+        td_vals = material_TDs.index_select(0, indices_rem)
+        col_vals = material_colors.index_select(0, indices_rem)
+        opac_vals = o + (
+            A * torch.log1p(k * (acc_thick[mask_remain] / td_vals))
+            + b * (acc_thick[mask_remain] / td_vals)
+        )
+        opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
+        comp[mask_remain] = comp[mask_remain] + (
+            (remaining[mask_remain] * opac_vals).unsqueeze(-1) * col_vals
+        )
+        remaining[mask_remain] = remaining[mask_remain] * (1 - opac_vals)
+        cur_mat[mask_remain] = -1
+        acc_thick[mask_remain] = 0.0
+
+    comp = comp + remaining.unsqueeze(-1) * background
+    return comp * 255.0
+
+
 @torch.jit.script
 def adaptive_round(
     x: torch.Tensor, tau: float, high_tau: float, low_tau: float, temp: float

{autoforge-1.4.2 → autoforge-1.5.0}/src/autoforge/Helper/OutputHelper.py

@@ -33,7 +33,7 @@ def extract_filament_swaps(disc_global, disc_height_image, background_layers):
         current = int(disc_global[i])
         # If this is the first layer or the material changes from the previous layer…
         if current != prev:
-            slider = (i + background_layers) - 1
+            slider = i + background_layers
             slider_values.append(slider)
             filament_indices.append(prev)
         prev = current
@@ -118,7 +118,7 @@ def generate_project_file(
         },
     )
     # add black to slider value
-    slider_values.insert(0, (args.background_height // args.layer_height) - 1)
+    slider_values.insert(0, (args.background_height // args.layer_height))
 
     # reverse order of filament set
     filament_set = filament_set[::-1]
@@ -186,11 +186,9 @@ def generate_project_file(
         json.dump(project_data, f, indent=4)
 
 
-import numpy as np
-import struct
-
-
-def generate_stl(height_map, filename, background_height, maximum_x_y_size, alpha_mask=None):
+def generate_stl(
+    height_map, filename, background_height, maximum_x_y_size, alpha_mask=None
+):
     """
     Generate a binary STL file from a height map with an optional alpha mask.
     If alpha_mask is provided, vertices where alpha < 128 are omitted.
@@ -213,7 +211,7 @@ def generate_stl(height_map, filename, background_height, maximum_x_y_size, alph
     if alpha_mask is None:
         valid_mask = np.ones((H, W), dtype=bool)
     else:
-        valid_mask = (alpha_mask >= 128)
+        valid_mask = alpha_mask >= 128
 
     # --- Create vertices for the top surface ---
     # We compute a grid of vertices (even for masked-out regions) but later only use those that are valid.
@@ -250,7 +248,12 @@ def generate_stl(height_map, filename, background_height, maximum_x_y_size, alph
     top_triangles_indices = []  # Each element is a tuple of 3 (i,j) indices.
     for i in range(H - 1):
         for j in range(W - 1):
-            if valid_mask[i, j] and valid_mask[i, j + 1] and valid_mask[i + 1, j + 1] and valid_mask[i + 1, j]:
+            if (
+                valid_mask[i, j]
+                and valid_mask[i, j + 1]
+                and valid_mask[i + 1, j + 1]
+                and valid_mask[i + 1, j]
+            ):
                 # Define the four corners of the quad.
                 v0 = top_vertices[i, j]
                 v1 = top_vertices[i, j + 1]
@@ -306,7 +309,12 @@ def generate_stl(height_map, filename, background_height, maximum_x_y_size, alph
     # The ordering is reversed (relative to the top face) so that the normals face downward.
     for i in range(H - 1):
         for j in range(W - 1):
-            if valid_mask[i, j] and valid_mask[i, j + 1] and valid_mask[i + 1, j + 1] and valid_mask[i + 1, j]:
+            if (
+                valid_mask[i, j]
+                and valid_mask[i, j + 1]
+                and valid_mask[i + 1, j + 1]
+                and valid_mask[i + 1, j]
+            ):
                 v0 = bottom_vertices[i, j]
                 v1 = bottom_vertices[i, j + 1]
                 v2 = bottom_vertices[i + 1, j + 1]
@@ -383,7 +391,7 @@ def generate_swap_instructions(
     instructions.append("Start with your background color")
     for i in range(0, L):
         if i == 0 or int(discrete_global[i]) != int(discrete_global[i - 1]):
-            ie = i
+            ie = i + 1
             instructions.append(
                 f"At layer #{ie + background_layers} ({(ie * h) + background_height:.2f}mm) swap to {material_names[int(discrete_global[i])]}"
             )

{autoforge-1.4.2 → autoforge-1.5.0}/src/autoforge/Helper/PruningHelper.py

@@ -1,7 +1,7 @@
 import torch
 from tqdm import tqdm
 
-from autoforge.Helper.OptimizerHelper import composite_image, discretize_solution
+from autoforge.Helper.OptimizerHelper import discretize_solution, composite_image_disc
 from autoforge.Loss.LossFunctions import compute_loss
 
 
@@ -55,7 +55,7 @@ def prune_num_colors(
         )
 
         with torch.no_grad():
-            out_im = composite_image(
+            out_im = composite_image_disc(
                 pixel_height_logits,
                 logits_for_disc,
                 tau_for_comp,
@@ -65,7 +65,6 @@ def prune_num_colors(
                 material_colors,
                 material_TDs,
                 background,
-                mode="discrete",
                 rng_seed=rng_seed,
             )
 
@@ -148,7 +147,7 @@ def prune_num_swaps(
         )
 
         with torch.no_grad():
-            out_im = composite_image(
+            out_im = composite_image_disc(
                 pixel_height_logits,
                 logits_for_disc,
                 tau_for_comp,
@@ -158,7 +157,6 @@ def prune_num_swaps(
                 material_colors,
                 material_TDs,
                 background,
-                mode="discrete",
                 rng_seed=rng_seed,
             )
         return compute_loss(
@@ -413,7 +411,7 @@ def prune_redundant_layers(
     current_max_layers = current_params["global_logits"].shape[0]
 
     # Compute baseline composite and loss.
-    comp = composite_image(
+    comp = composite_image_disc(
         current_params["pixel_height_logits"],
         current_params["global_logits"],
         final_tau,
@@ -423,7 +421,6 @@ def prune_redundant_layers(
         material_colors,
         material_TDs,
         background,
-        mode="discrete",
         rng_seed=rng_seed,
     )
     disc_global, _ = discretize_solution(
@@ -458,7 +455,7 @@ def prune_redundant_layers(
             candidate_params, candidate_max_layers = remove_layer_from_solution(
                 current_params, layer, final_tau, h, current_max_layers, rng_seed
             )
-            candidate_comp = composite_image(
+            candidate_comp = composite_image_disc(
                 candidate_params["pixel_height_logits"],
                 candidate_params["global_logits"],
                 final_tau,
@@ -468,7 +465,6 @@ def prune_redundant_layers(
                 material_colors,
                 material_TDs,
                 background,
-                mode="discrete",
                 rng_seed=rng_seed,
             )
             candidate_disc_global, _ = discretize_solution(

{autoforge-1.4.2 → autoforge-1.5.0}/src/autoforge/Loss/LossFunctions.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn.functional as F
 
 from autoforge.Helper.ImageHelper import srgb_to_lab, increase_saturation
-from autoforge.Helper.OptimizerHelper import composite_image
+from autoforge.Helper.OptimizerHelper import composite_image_cont
 
 
 def loss_fn(
@@ -22,7 +22,7 @@ def loss_fn(
     Full forward pass for continuous assignment:
     composite, then compute unified loss on (global_logits).
     """
-    comp = composite_image(
+    comp = composite_image_cont(
         params["pixel_height_logits"],
         params["global_logits"],
         tau_height,
@@ -32,7 +32,6 @@ def loss_fn(
         material_colors,
         material_TDs,
         background,
-        mode="continuous",
     )
     return compute_loss(
         material_assignment=params["global_logits"],

{autoforge-1.4.2 → autoforge-1.5.0}/src/autoforge/Modules/Optimizer.py

@@ -5,11 +5,15 @@ import random
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
-import torch.optim as optim
 from tqdm import tqdm
 from torch.utils.tensorboard import SummaryWriter
 
-from autoforge.Helper.OptimizerHelper import composite_image, discretize_solution
+from autoforge.Helper.CAdamW import CAdamW
+from autoforge.Helper.OptimizerHelper import (
+    discretize_solution,
+    composite_image_cont,
+    composite_image_disc,
+)
 from autoforge.Helper.PruningHelper import prune_colors_and_swaps, find_color_bands
 from autoforge.Loss.LossFunctions import loss_fn, compute_loss
 
@@ -64,7 +68,7 @@ class FilamentOptimizer:
         # Initialize TensorBoard writer
         if args.tensorboard:
             if args.run_name:
-                self.writer = SummaryWriter(log_dir=f'runs/{args.run_name}')
+                self.writer = SummaryWriter(log_dir=f"runs/{args.run_name}")
             else:
                 self.writer = SummaryWriter()
         else:
@@ -99,9 +103,10 @@ class FilamentOptimizer:
         )
 
         # Initialize optimizer
-        self.optimizer = optim.Adam(
+        self.optimizer = CAdamW(
             [self.params["pixel_height_logits"], self.params["global_logits"]],
             lr=self.learning_rate,
+            weight_decay=1e-2,
         )
 
         # Setup best discrete solution tracking
@@ -193,7 +198,9 @@ class FilamentOptimizer:
 
         return loss
 
-    def log_to_tensorboard(self, interval: int = 100, namespace: str = "", step: int = None):
+    def log_to_tensorboard(
+        self, interval: int = 100, namespace: str = "", step: int = None
+    ):
         """
         Log metrics and images to TensorBoard.
 
@@ -211,22 +218,26 @@ class FilamentOptimizer:
         steps = step if step is not None else self.num_steps_done
 
         # Log metrics
-        self.writer.add_scalar(f'Loss/{prefix}best_discrete', self.best_discrete_loss, steps)
-        self.writer.add_scalar(f'Loss/{prefix}best_swaps', self.best_swaps, steps)
+        self.writer.add_scalar(
+            f"Loss/{prefix}best_discrete", self.best_discrete_loss, steps
+        )
+        self.writer.add_scalar(f"Loss/{prefix}best_swaps", self.best_swaps, steps)
 
         tau_height, tau_global = self._get_tau()
 
         # Metrics that are only relevant for the main optimization loop
         if not prefix:
-            self.writer.add_scalar('Params/tau_height', tau_height, steps)
-            self.writer.add_scalar('Params/tau_global', tau_global, steps)
-            self.writer.add_scalar('Params/lr', self.optimizer.param_groups[0]['lr'], steps)
-            self.writer.add_scalar('Loss/train', self.loss, steps)
+            self.writer.add_scalar("Params/tau_height", tau_height, steps)
+            self.writer.add_scalar("Params/tau_global", tau_global, steps)
+            self.writer.add_scalar(
+                "Params/lr", self.optimizer.param_groups[0]["lr"], steps
+            )
+            self.writer.add_scalar("Loss/train", self.loss, steps)
 
         # Log images periodically
         if (steps + 1) % interval == 0:
             with torch.no_grad():
-                comp_img = composite_image(
+                comp_img = composite_image_cont(
                     self.params["pixel_height_logits"],
                     self.params["global_logits"],
                     tau_height,
@@ -235,9 +246,13 @@ class FilamentOptimizer:
                     self.max_layers,
                     self.material_colors,
                     self.material_TDs,
-                    self.background
+                    self.background,
+                )
+                self.writer.add_images(
+                    f"Current Output/{prefix}composite",
+                    comp_img.permute(2, 0, 1).unsqueeze(0) / 255.0,
+                    steps,
                 )
-                self.writer.add_images(f'Current Output/{prefix}composite', comp_img.permute(2, 0, 1).unsqueeze(0) / 255.0, steps)
 
     def visualize(self, interval: int = 25):
         """
@@ -255,7 +270,7 @@ class FilamentOptimizer:
 
         with torch.no_grad():
             tau_h, tau_g = self._get_tau()
-            comp = composite_image(
+            comp = composite_image_cont(
                 self.params["pixel_height_logits"],
                 self.params["global_logits"],
                 tau_h,
@@ -265,14 +280,13 @@ class FilamentOptimizer:
                 self.material_colors,
                 self.material_TDs,
                 self.background,
-                mode="continuous",
             )
         comp_np = np.clip(comp.cpu().detach().numpy(), 0, 255).astype(np.uint8)
         self.current_comp_ax.set_data(comp_np)
 
         if self.best_params is not None:
             with torch.no_grad():
-                best_comp = composite_image(
+                best_comp = composite_image_disc(
                     self.best_params["pixel_height_logits"],
                     self.best_params["global_logits"],
                     self.final_tau,
@@ -282,7 +296,6 @@ class FilamentOptimizer:
                     self.material_colors,
                     self.material_TDs,
                     self.background,
-                    mode="discrete",
                     rng_seed=self.best_seed,
                 )
             best_comp_np = np.clip(best_comp.cpu().detach().numpy(), 0, 255).astype(
@@ -434,7 +447,7 @@ class FilamentOptimizer:
 
             # 2) Compute discrete-mode composite
             with torch.no_grad():
-                comp_disc = composite_image(
+                comp_disc = composite_image_disc(
                     self.params["pixel_height_logits"],
                     self.params["global_logits"],
                     self.final_tau,
@@ -444,7 +457,6 @@ class FilamentOptimizer:
                     self.material_colors,
                     self.material_TDs,
                     self.background,
-                    mode="discrete",  # or "discrete" if you prefer fully discrete
                     rng_seed=seed,
                 )
 
@@ -487,7 +499,7 @@ class FilamentOptimizer:
             disc_global, disc_height_image = discretize_solution(
                 self.best_params, self.final_tau, self.h, self.max_layers, rng_seed=seed
             )
-            comp_disc = composite_image(
+            comp_disc = composite_image_disc(
                 self.best_params["pixel_height_logits"],
                 self.best_params["global_logits"],
                 self.final_tau,
@@ -497,7 +509,6 @@ class FilamentOptimizer:
                 self.material_colors,
                 self.material_TDs,
                 self.background,
-                mode="discrete",
                 rng_seed=seed,
             )
             current_disc_loss = compute_loss(

{autoforge-1.4.2 → autoforge-1.5.0}/src/autoforge/auto_forge.py

@@ -10,9 +10,9 @@ from autoforge.Helper.FilamentHelper import hex_to_rgb, load_materials
 from autoforge.Helper.ImageHelper import resize_image, resize_image_exact
 from autoforge.Helper.OptimizerHelper import (
     init_height_map,
-    composite_image,
     discretize_solution,
     init_height_map_depth_color_adjusted,
+    composite_image_disc,
 )
 from autoforge.Helper.OutputHelper import (
     generate_stl,
@@ -207,6 +207,9 @@ def main():
         "Background must be multiple of layer_height"
     )
 
+    assert os.path.exists(args.input_image), "Input image not found"
+    assert os.path.exists(args.csv_file), "CSV file not found"
+
     random_seed = args.random_seed
     if random_seed == 0:
         random_seed = int(time.time())
@@ -387,7 +390,7 @@ def main():
 
     print("Done. Saving outputs...")
     # Save Image
-    comp_disc = composite_image(
+    comp_disc = composite_image_disc(
         params["pixel_height_logits"],
         params["global_logits"],
         args.final_tau,
@@ -397,7 +400,6 @@ def main():
         material_colors,
         material_TDs,
         background,
-        mode="discrete",  # or "discrete" if you prefer fully discrete
         rng_seed=optimizer.best_seed,
     )