autoforge 1.4.2__tar.gz → 1.5.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.4.2
+Version: 1.5.1
 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
@@ -25,16 +25,6 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
 
 **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
 
-
-## Important Information
-
-We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
-This can have some unforeseen consequences, so please report any bugs you find. 
-
-This release should give better gpu support for more users, especially on windows. \
-If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
-Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
-
 ## Example
 All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
 <div style="display: flex; justify-content: center; gap: 20px;">
@@ -46,9 +36,9 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
   </div>
   <div style="text-align: center;">
     <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
-    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="200" />
+    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.png" width="200" />
     <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
-    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.jpg" width="200" />
+    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.png" width="200" />
   </div>
   <div style="text-align: center;">
     <h3>Imported Hueforge View</h3>
@@ -80,6 +70,9 @@ To install AutoForge, simply install the current version from PyPI:
    pip install autoforge
 ```
 
+If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
+Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
+
 ## Usage
 
 The script is run from the command line and accepts several arguments. Below is an example command:
@@ -125,7 +118,7 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
 
 ### Experimental Deph Anything V2 parameters
-I got a request to add a initializing function that takes the original deph of the image into account when initializing. \
+I got a request to add an initializing function that takes the original depth of the image into account when initializing. \
 In theory this should give you a nice background/foreground separation, but in practice it can degrade the quality of the output. \
 In addition you need the VRAM to run the depth model. \
 I currently don't recommend it as the output is not as good as without it, but I will leave it in for now. \
@@ -172,7 +165,7 @@ The software is provided as-is and comes with no warranty or guarantee of suppor
 ## Acknowledgements
 
 First and foremost:
-- [Hueforge](https://shop.thehueforge.com/) for providing the filament data and inspiration for this project.
+- [Hueforge](https://shop.thehueforge.com/) for providing the inspiration for this project.
 Without it, this project would not have been possible.
 
 AutoForge makes use of several open source libraries:

{autoforge-1.4.2 → autoforge-1.5.1}/README.md

@@ -4,16 +4,6 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
 
 **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
 
-
-## Important Information
-
-We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
-This can have some unforeseen consequences, so please report any bugs you find. 
-
-This release should give better gpu support for more users, especially on windows. \
-If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
-Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
-
 ## Example
 All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
 <div style="display: flex; justify-content: center; gap: 20px;">
@@ -25,9 +15,9 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
   </div>
   <div style="text-align: center;">
     <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
-    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="200" />
+    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.png" width="200" />
     <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
-    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.jpg" width="200" />
+    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.png" width="200" />
   </div>
   <div style="text-align: center;">
     <h3>Imported Hueforge View</h3>
@@ -59,6 +49,9 @@ To install AutoForge, simply install the current version from PyPI:
    pip install autoforge
 ```
 
+If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
+Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
+
 ## Usage
 
 The script is run from the command line and accepts several arguments. Below is an example command:
@@ -104,7 +97,7 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
 
 ### Experimental Deph Anything V2 parameters
-I got a request to add a initializing function that takes the original deph of the image into account when initializing. \
+I got a request to add an initializing function that takes the original depth of the image into account when initializing. \
 In theory this should give you a nice background/foreground separation, but in practice it can degrade the quality of the output. \
 In addition you need the VRAM to run the depth model. \
 I currently don't recommend it as the output is not as good as without it, but I will leave it in for now. \
@@ -151,7 +144,7 @@ The software is provided as-is and comes with no warranty or guarantee of suppor
 ## Acknowledgements
 
 First and foremost:
-- [Hueforge](https://shop.thehueforge.com/) for providing the filament data and inspiration for this project.
+- [Hueforge](https://shop.thehueforge.com/) for providing the inspiration for this project.
 Without it, this project would not have been possible.
 
 AutoForge makes use of several open source libraries:

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.4.2
+Version: 1.5.1
 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
@@ -25,16 +25,6 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
 
 **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
 
-
-## Important Information
-
-We recently switched from JAX to PyTorch, to allow for a more streamlined development process. \
-This can have some unforeseen consequences, so please report any bugs you find. 
-
-This release should give better gpu support for more users, especially on windows. \
-If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
-Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
-
 ## Example
 All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
 <div style="display: flex; justify-content: center; gap: 20px;">
@@ -46,9 +36,9 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
   </div>
   <div style="text-align: center;">
     <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
-    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.jpg" width="200" />
+    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.png" width="200" />
     <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
-    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.jpg" width="200" />
+    <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.png" width="200" />
   </div>
   <div style="text-align: center;">
     <h3>Imported Hueforge View</h3>
@@ -80,6 +70,9 @@ To install AutoForge, simply install the current version from PyPI:
    pip install autoforge
 ```
 
+If you have problems running the code on your gpu, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
+Both CUDA and ROCm are supported, but you need to install the correct version of pytorch for your system.
+
 ## Usage
 
 The script is run from the command line and accepts several arguments. Below is an example command:
@@ -125,7 +118,7 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
 - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
 
 ### Experimental Deph Anything V2 parameters
-I got a request to add a initializing function that takes the original deph of the image into account when initializing. \
+I got a request to add an initializing function that takes the original depth of the image into account when initializing. \
 In theory this should give you a nice background/foreground separation, but in practice it can degrade the quality of the output. \
 In addition you need the VRAM to run the depth model. \
 I currently don't recommend it as the output is not as good as without it, but I will leave it in for now. \
@@ -172,7 +165,7 @@ The software is provided as-is and comes with no warranty or guarantee of suppor
 ## Acknowledgements
 
 First and foremost:
-- [Hueforge](https://shop.thehueforge.com/) for providing the filament data and inspiration for this project.
+- [Hueforge](https://shop.thehueforge.com/) for providing the inspiration for this project.
 Without it, this project would not have been possible.
 
 AutoForge makes use of several open source libraries:

{autoforge-1.4.2 → autoforge-1.5.1}/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.1/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.1}/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,162 @@ def composite_image(
     return comp * 255.0
 
 
+@torch.jit.script
+def composite_image_disc(
+    pixel_height_logits: torch.Tensor,  # [H,W]
+    global_logits: torch.Tensor,  # [max_layers, n_materials]
+    tau_height: float,
+    tau_global: float,
+    h: float,
+    max_layers: int,
+    material_colors: torch.Tensor,  # [n_materials, 3]
+    material_TDs: torch.Tensor,  # [n_materials]
+    background: torch.Tensor,  # [3]
+    rng_seed: int = -1,
+) -> torch.Tensor:
+    """
+    An optimized rewrite of the original compositing code.
+
+    The main changes:
+      - We combine the old (mask_diff) and (mask_not_printed) logic
+        into a single step that composites any pixel which is 'done'
+        with its current material (either because it's not printing
+        this layer or because the layer material is different).
+      - Then we update (or start) new segments for any pixel that
+        actually prints this layer and whose material matches or
+        doesn't match.
+      - Fewer conditionals, fewer index_select calls, and no repeated
+        if mask.any(): checks.
+    """
+
+    # Determine how many discrete layers each pixel prints
+    pixel_height = (max_layers * h) * torch.sigmoid(pixel_height_logits)
+    continuous_layers = pixel_height / h
+
+    # Example "adaptive rounding" approach
+    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]
+
+    # Precompute the chosen material for each layer
+    # (One material per layer for the entire image.)
+    if rng_seed >= 0:
+        # Deterministic sampling for each layer
+        new_mats_list = []
+        for layer_idx in range(max_layers):
+            p_i = deterministic_gumbel_softmax(
+                global_logits[layer_idx],
+                tau_global,
+                hard=True,
+                rng_seed=(rng_seed + layer_idx),
+            )
+            mat_i = torch.argmax(p_i, dim=0).to(torch.int32)
+            new_mats_list.append(mat_i)
+        new_mats = torch.stack(new_mats_list, dim=0)  # [max_layers]
+    else:
+        # Standard (random) Gumbel softmax
+        p_all = F.gumbel_softmax(global_logits, tau_global, hard=True, dim=1)  # [L, M]
+        new_mats = torch.argmax(p_all, dim=1).to(torch.int32)  # [max_layers]
+
+
+    H, W = pixel_height.shape
+    device = pixel_height.device
+
+    comp = torch.zeros(H, W, 3, dtype=torch.float32, device=device)
+    remaining = torch.ones(H, W, dtype=torch.float32, device=device)
+
+    # Tracks the material currently being accumulated (or -1 for none)
+    cur_mat = -torch.ones((H, W), dtype=torch.int32, device=device)
+    # Tracks the accumulated thickness for the current segment
+    acc_thick = torch.zeros((H, W), dtype=torch.float32, device=device)
+
+    # Opacity function parameters
+    o = 0.10868816
+    A = 0.3077416
+    k = 76.928215
+    b = 2.2291653
+
+    # Main compositing loop: top to bottom
+    for layer_idx in range(max_layers - 1, -1, -1):
+        # Material chosen for this layer
+        layer_mat = new_mats[layer_idx]  # a single int
+        # Which pixels actually print this layer?
+        p_print = discrete_layers > layer_idx  # [H,W] bool
+
+        # --------------------------------------------------
+        # (A) Composite any "finished" segment, i.e. the pixel
+        #     is *done* with its old segment if:
+        #        (1) we already had a segment (cur_mat != -1), AND
+        #        (2) either the pixel is not printing this layer,
+        #            or the pixel's new layer material != cur_mat
+        # --------------------------------------------------
+        mask_done = (cur_mat != -1) & (~p_print | (cur_mat != layer_mat))
+        if mask_done.any():
+            old_inds = cur_mat[mask_done].to(torch.int64)
+            td_vals = material_TDs[old_inds]  # gather thickness densities
+            col_vals = material_colors[old_inds]  # gather old material colors
+
+            thick_ratio = acc_thick[mask_done] / td_vals
+            opac_vals = o + (A * torch.log1p(k * thick_ratio) + b * thick_ratio)
+            opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
+
+            # composite old segment
+            comp[mask_done] += (remaining[mask_done] * opac_vals).unsqueeze(
+                -1
+            ) * col_vals
+            remaining[mask_done] *= 1.0 - opac_vals
+
+            # Reset old segment
+            cur_mat[mask_done] = -1
+            acc_thick[mask_done] = 0.0
+
+        # --------------------------------------------------
+        # (B) For pixels that do print this layer:
+        #       - start a new segment if cur_mat == -1
+        #       - accumulate thickness if cur_mat == layer_mat
+        # --------------------------------------------------
+        eff_thick = p_print.to(torch.float32) * h
+
+        # Start new segment
+        mask_new = p_print & (cur_mat == -1)
+        if mask_new.any():
+            cur_mat[mask_new] = layer_mat
+            acc_thick[mask_new] = eff_thick[mask_new]
+
+        # Accumulate thickness
+        mask_same = p_print & (cur_mat == layer_mat)
+        if mask_same.any():
+            acc_thick[mask_same] += eff_thick[mask_same]
+
+
+
+    # After the loop, composite any still-pending segments
+    mask_remain = cur_mat != -1
+    if mask_remain.any():
+        old_inds = cur_mat[mask_remain].to(torch.int64)
+        td_vals = material_TDs[old_inds]
+        col_vals = material_colors[old_inds]
+
+        thick_ratio = acc_thick[mask_remain] / td_vals
+        opac_vals = o + (A * torch.log1p(k * thick_ratio) + b * thick_ratio)
+        opac_vals = torch.clamp(opac_vals, 0.0, 1.0)
+
+        comp[mask_remain] += (remaining[mask_remain] * opac_vals).unsqueeze(
+            -1
+        ) * col_vals
+        remaining[mask_remain] *= 1.0 - opac_vals
+
+    # -------------------------------------------------------------------------
+    # Composite background
+    # -------------------------------------------------------------------------
+    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.1}/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.1}/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.1}/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.1}/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,7 +103,7 @@ class FilamentOptimizer:
         )
 
         # Initialize optimizer
-        self.optimizer = optim.Adam(
+        self.optimizer = CAdamW(
             [self.params["pixel_height_logits"], self.params["global_logits"]],
             lr=self.learning_rate,
         )
@@ -193,7 +197,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 +217,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 +245,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 +269,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 +279,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 +295,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 +446,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 +456,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 +498,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 +508,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.1}/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,
     )