autoforge 1.5.0__tar.gz → 1.5.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.5.0
+Version: 1.5.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
@@ -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.5.0 → autoforge-1.5.2}/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.5.0 → autoforge-1.5.2}/pyproject.toml

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: AutoForge
-Version: 1.5.0
+Version: 1.5.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
@@ -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.5.0 → autoforge-1.5.2}/src/autoforge/Helper/OptimizerHelper.py

@@ -20,6 +20,16 @@ def adaptive_round(
 ) -> torch.Tensor:
     """
     Smooth rounding based on temperature 'tau'.
+
+    Args:
+        x (torch.Tensor): The input tensor to be rounded.
+        tau (float): The current temperature parameter.
+        high_tau (float): The high threshold for the temperature.
+        low_tau (float): The low threshold for the temperature.
+        temp (float): The temperature parameter for the sigmoid function.
+
+    Returns:
+        torch.Tensor: The rounded tensor.
     """
     if tau <= low_tau:
         return torch.round(x)
@@ -40,6 +50,16 @@ def adaptive_round(
 # A deterministic random generator that mimics torch.rand_like.
 @torch.jit.script
 def deterministic_rand_like(tensor: torch.Tensor, seed: int) -> torch.Tensor:
+    """
+    Generate a deterministic random tensor that mimics torch.rand_like.
+
+    Args:
+        tensor (torch.Tensor): The input tensor whose shape and device will be used.
+        seed (int): The seed for the deterministic random generator.
+
+    Returns:
+        torch.Tensor: A tensor with the same shape as the input tensor, filled with deterministic random values.
+    """
     # Compute the total number of elements.
     n: int = 1
     for d in tensor.shape:
@@ -59,6 +79,18 @@ def deterministic_rand_like(tensor: torch.Tensor, seed: int) -> torch.Tensor:
 def deterministic_gumbel_softmax(
     logits: torch.Tensor, tau: float, hard: bool, rng_seed: int
 ) -> torch.Tensor:
+    """
+    Apply the Gumbel-Softmax trick in a deterministic manner using a fixed random seed.
+
+    Args:
+        logits (torch.Tensor): The input logits tensor.
+        tau (float): The temperature parameter for the Gumbel-Softmax.
+        hard (bool): If True, the output will be one-hot encoded.
+        rng_seed (int): The seed for the deterministic random generator.
+
+    Returns:
+        torch.Tensor: The resulting tensor after applying the Gumbel-Softmax trick.
+    """
     eps: float = 1e-20
     # Instead of torch.rand_like(..., generator=...), use our deterministic_rand_like.
     U = deterministic_rand_like(logits, rng_seed)
@@ -83,218 +115,341 @@ def composite_image_cont(
     tau_global: float,
     h: float,
     max_layers: int,
-    material_colors: torch.Tensor,
-    material_TDs: torch.Tensor,
-    background: torch.Tensor,
+    material_colors: torch.Tensor,  # [n_materials, 3]
+    material_TDs: torch.Tensor,  # [n_materials]
+    background: torch.Tensor,  # [3]
 ) -> torch.Tensor:
     """
     Continuous compositing over all pixels.
     Uses Gumbel softmax with either hard or soft sampling depending on tau_global.
+
+    Args:
+        pixel_height_logits (torch.Tensor): Logits for pixel heights, shape [H, W].
+        global_logits (torch.Tensor): Logits for global material assignment, shape [max_layers, n_materials].
+        tau_height (float): Temperature parameter for height rounding.
+        tau_global (float): Temperature parameter for global material assignment.
+        h (float): Height of each layer.
+        max_layers (int): Maximum number of layers.
+        material_colors (torch.Tensor): Tensor of material colors, shape [n_materials, 3].
+        material_TDs (torch.Tensor): Tensor of material transmission/opacity parameters, shape [n_materials].
+        background (torch.Tensor): Background color tensor, shape [3].
+
+    Returns:
+        torch.Tensor: Composite image tensor, shape [H, W, 3].
     """
+    # Compute continuous and discrete layers for each pixel
     pixel_height = (max_layers * h) * torch.sigmoid(pixel_height_logits)
     continuous_layers = pixel_height / h
 
+    # adaptive rounding
     adaptive_layers = adaptive_round(
         continuous_layers, tau_height, high_tau=0.1, low_tau=0.01, temp=0.1
     )
+
+    # "Straight-through" trick to keep gradient from adaptive rounding
     discrete_layers_temp = torch.round(continuous_layers)
     discrete_layers = (
         discrete_layers_temp + (adaptive_layers - discrete_layers_temp).detach()
-    ).to(torch.int32)
-
+    ).to(torch.int32)  # shape [H, W]
+
+    # Precompute global material assignments for each layer
+    # Use a single Gumbel Softmax call for all layers.
+    hard_flag = tau_global < 1e-3
+    # global_logits: [max_layers, n_materials]
+    p = F.gumbel_softmax(
+        global_logits, tau_global, hard=hard_flag, dim=1
+    )  # [L, n_materials]
+
+    # Compute color_i and TD_i for each layer in one go.
+    # color_i -> [max_layers, 3], TD_i -> [max_layers]
+    layer_colors = p @ material_colors  # [L, 3]
+    layer_TDs = p @ material_TDs  # [L]
+    layer_TDs.clamp_(1e-8, 1e8)
+
+    # Prepare output and compositing variables
     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
+    # Opacity function parameters
     o = -1.2416557e-02
     A = 9.6407950e-01
     k = 3.4103447e01
     b = -4.1554203e00
 
+    # Composite from top layer (max_layers-1) down to 0
     for i in range(max_layers):
         layer_idx = max_layers - 1 - i
-        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)
+        # Mask of which pixels actually print this layer:
+        p_print = (discrete_layers > layer_idx).to(dtype=comp.dtype)
 
-        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)
+        # Effective thickness for those pixels
+        eff_thick = p_print * h
 
+        # Opacity based on thickness and material TD
+        TD_i = layer_TDs[layer_idx]
         opac = o + (A * torch.log1p(k * (eff_thick / TD_i)) + b * (eff_thick / TD_i))
         opac = torch.clamp(opac, 0.0, 1.0)
 
-        comp = comp + ((remaining * opac).unsqueeze(-1) * color_i)
-        remaining = remaining * (1 - opac)
+        # Add contribution to comp
+        color_i = layer_colors[layer_idx]  # shape [3]
+        # We need to broadcast opac: [H, W] -> [H, W, 1]
+        comp = comp + (remaining * opac).unsqueeze(-1) * color_i
+
+        # Update the remaining factor
+        remaining = remaining * (1.0 - opac)
 
+    # Add background contribution
     comp = comp + remaining.unsqueeze(-1) * background
+
+    # 6. Scale to [0, 255] range
     return comp * 255.0
 
 
 @torch.jit.script
 def composite_image_disc(
-    pixel_height_logits: torch.Tensor,
-    global_logits: torch.Tensor,
+    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,
-    material_TDs: torch.Tensor,
-    background: torch.Tensor,
+    material_colors: torch.Tensor,  # [n_materials, 3]
+    material_TDs: torch.Tensor,  # [n_materials]
+    background: torch.Tensor,  # [3]
     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.
+    Perform discrete compositing over all pixels.
+
+    Args:
+        pixel_height_logits (torch.Tensor): Logits for pixel heights, shape [H, W].
+        global_logits (torch.Tensor): Logits for global material assignment, shape [max_layers, n_materials].
+        tau_height (float): Temperature parameter for height rounding.
+        tau_global (float): Temperature parameter for global material assignment.
+        h (float): Height of each layer.
+        max_layers (int): Maximum number of layers.
+        material_colors (torch.Tensor): Tensor of material colors, shape [n_materials, 3].
+        material_TDs (torch.Tensor): Tensor of material transmission/opacity parameters, shape [n_materials].
+        background (torch.Tensor): Background color tensor, shape [3].
+        rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
+
+    Returns:
+        torch.Tensor: Composite image tensor, shape [H, W, 3].
     """
-    # Compute pixel height and determine discrete layers.
+
+    # -------------------------------------------------------------------------
+    # 1) Compute discrete per-pixel layer counts (discrete_layers).
+    # -------------------------------------------------------------------------
+    #  pixel_height ~ [0, max_layers*h]
     pixel_height = (max_layers * h) * torch.sigmoid(pixel_height_logits)
+
+    #  continuous_layers ~ [0, max_layers]
     continuous_layers = pixel_height / h
 
+    #  Use your "adaptive rounding" trick if desired:
     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:
+    ).to(torch.int32)  # [H,W]
+
+    # -------------------------------------------------------------------------
+    # 2) Pick a single global material per layer, either deterministically
+    #    or via gumbel softmax, as in original code.
+    # -------------------------------------------------------------------------
+    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, 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)
+                global_logits[layer_idx],
+                tau_global,
+                hard=True,
+                rng_seed=(rng_seed + layer_idx),
             )
-            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
+            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)
+
+    # Current material index for each pixel, or -1 for none
+    cur_mat = -torch.ones((H, W), dtype=torch.int32, device=device)
+
+    # 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):
+        # layer_mat is the global material chosen for this layer (int32 scalar).
+        layer_mat = new_mats[layer_idx]  # shape []
+
+        # Which pixels actually print on this layer?
+        # p_print = (discrete_layers > layer_idx)
+        p_print = discrete_layers.gt(layer_idx)  # bool
+
+        # ---------------------------------------------------------------------
+        # (A) "Finish" any existing segments that are now 'done'.
+        #
+        # A segment is done if:
+        #   1) cur_mat != -1, i.e. the pixel had an ongoing segment
+        #   2) EITHER
+        #       - the pixel does not print now (~p_print),
+        #       - OR the new layer material differs (cur_mat != layer_mat).
+        # ---------------------------------------------------------------------
+        mask_done = (cur_mat.ne(-1)) & ((~p_print) | (cur_mat.ne(layer_mat)))
+
+        # Convert to float for multiplications
+        mask_done_f = mask_done.to(torch.float32)
+
+        # Gather thickness densities & colors for the old segment
+        # We'll clamp cur_mat so -1 becomes 0 (doesn't matter since we multiply by 0).
+        old_inds_clamped = torch.clamp(cur_mat, min=0)
+        td_vals = material_TDs[old_inds_clamped]  # [H, W]
+        col_vals = material_colors[old_inds_clamped]  # [H, W, 3]
+
+        # Compute alpha from accumulated thickness
+        thick_ratio = acc_thick / td_vals
+        opac_vals = o + (A * torch.log1p(k * thick_ratio) + b * thick_ratio)
+        opac_vals = torch.clamp(opac_vals, 0.0, 1.0)  # [H, W]
+
+        # Compositing the old segment:
+        #   comp += mask_done_f * remaining * opac_vals * col_vals
+        #   remaining *= (1 - mask_done_f * opac_vals) ...
+        # but we have to broadcast for color:
+        comp_add = (mask_done_f * remaining * opac_vals).unsqueeze(-1) * col_vals
+        comp = comp + comp_add
+        remaining = remaining - (mask_done_f * remaining * opac_vals)
+
+        # Reset old segment where mask_done is True
+        #   cur_mat = -1,  acc_thick = 0
+        # We'll do it by `torch.where(mask, val_if_true, val_if_false)`
+        cur_mat = torch.where(mask_done, torch.full_like(cur_mat, -1), cur_mat)
+        acc_thick = torch.where(mask_done, torch.zeros_like(acc_thick), acc_thick)
+
+        # ---------------------------------------------------------------------
+        # (B) For pixels that 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
+
+        # (B1) Start new segment where cur_mat == -1
+        mask_new = p_print & (cur_mat.eq(-1))
+        mask_new_f = mask_new.to(torch.float32)
+
+        # Set cur_mat to layer_mat if mask_new is True
+        # (layer_mat is shape [], so it will broadcast)
+        cur_mat = torch.where(mask_new, layer_mat, cur_mat)
+
+        # We add thickness:
+        acc_thick = acc_thick + mask_new_f * eff_thick
+
+        # (B2) Accumulate thickness where cur_mat == layer_mat
+        # We do this in a second mask to avoid confusion, but you can combine.
+        mask_same = p_print & (cur_mat.eq(layer_mat))
+        acc_thick = acc_thick + (mask_same.to(torch.float32) * eff_thick)
+
+    # -------------------------------------------------------------------------
+    # 5) After the loop, composite any remaining segments (cur_mat != -1).
+    # -------------------------------------------------------------------------
+    mask_remain = cur_mat.ne(-1)
+    mask_remain_f = mask_remain.to(torch.float32)
+
+    old_inds_clamped = torch.clamp(cur_mat, min=0)
+    td_vals = material_TDs[old_inds_clamped]
+    col_vals = material_colors[old_inds_clamped]
+
+    thick_ratio = acc_thick / 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_add = (mask_remain_f * remaining * opac_vals).unsqueeze(-1) * col_vals
+    comp = comp + comp_add
+    remaining = remaining - (mask_remain_f * remaining * opac_vals)
+
+    # -------------------------------------------------------------------------
+    # 6) Composite background
+    # -------------------------------------------------------------------------
     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
+def composite_image_combined(
+    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:
     """
-    Smooth rounding based on temperature 'tau'.
+    Combine continuous and discrete compositing over all pixels.
+
+    Args:
+        pixel_height_logits (torch.Tensor): Logits for pixel heights, shape [H, W].
+        global_logits (torch.Tensor): Logits for global material assignment, shape [max_layers, n_materials].
+        tau_height (float): Temperature parameter for height rounding.
+        tau_global (float): Temperature parameter for global material assignment.
+        h (float): Height of each layer.
+        max_layers (int): Maximum number of layers.
+        material_colors (torch.Tensor): Tensor of material colors, shape [n_materials, 3].
+        material_TDs (torch.Tensor): Tensor of material transmission/opacity parameters, shape [n_materials].
+        background (torch.Tensor): Background color tensor, shape [3].
+        rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
+
+    Returns:
+        torch.Tensor: Composite image tensor, shape [H, W, 3].
     """
-    if tau <= low_tau:
-        return torch.round(x)
-    elif tau >= high_tau:
-        floor_val = torch.floor(x)
-        diff = x - floor_val
-        soft_round = floor_val + torch.sigmoid((diff - 0.5) / temp)
-        return soft_round
+    cont = composite_image_cont(
+        pixel_height_logits,
+        global_logits,
+        tau_height,
+        tau_global,
+        h,
+        max_layers,
+        material_colors,
+        material_TDs,
+        background,
+    )
+    if tau_global < 1.0:
+        disc = composite_image_disc(
+            pixel_height_logits,
+            global_logits,
+            tau_height,
+            tau_global,
+            h,
+            max_layers,
+            material_colors,
+            material_TDs,
+            background,
+            rng_seed,
+        )
+        return cont * tau_global + disc * (1 - tau_global)
     else:
-        ratio = (tau - low_tau) / (high_tau - low_tau)
-        hard_round = torch.round(x)
-        floor_val = torch.floor(x)
-        diff = x - floor_val
-        soft_round = floor_val + torch.sigmoid((diff - 0.5) / temp)
-        return ratio * soft_round + (1 - ratio) * hard_round
+        return cont
 
 
 def discretize_solution(
@@ -302,6 +457,18 @@ def discretize_solution(
 ):
     """
     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"]
@@ -328,10 +495,10 @@ def initialize_pixel_height_logits(target):
     Uses the formula: L = 0.299*R + 0.587*G + 0.114*B.
 
     Args:
-        target (jnp.ndarray): The target image array with shape (H, W, 3).
+        target (np.ndarray): The target image array with shape (H, W, 3).
 
     Returns:
-        jnp.ndarray: The initialized pixel height logits.
+        np.ndarray: The initialized pixel height logits.
     """
     # Compute normalized luminance in [0,1]
     normalized_lum = (
@@ -711,6 +878,10 @@ def init_height_map(target, max_layers, h, eps=1e-6, random_seed=None):
 
     Args:
         target (jnp.ndarray): The target image array with shape (H, W, 3).
+        max_layers (int): Maximum number of layers.
+        h (float): Height of each layer.
+        eps (float): Small constant to avoid division by zero.
+        random_seed (int): Random seed for reproducibility.
 
     Returns:
         jnp.ndarray: The initialized pixel height logits.

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

@@ -106,7 +106,6 @@ class FilamentOptimizer:
         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

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

@@ -20,7 +20,6 @@ from autoforge.Helper.OutputHelper import (
     generate_project_file,
 )
 from autoforge.Helper.PruningHelper import prune_redundant_layers
-from autoforge.Loss.PerceptionLoss import MultiLayerVGGPerceptualLoss
 from autoforge.Modules.Optimizer import FilamentOptimizer
 
 
@@ -282,7 +281,8 @@ def main():
     )  # shape (H', W') cause of weird opencv quirk
 
     # VGG Perceptual Loss
-    perception_loss_module = MultiLayerVGGPerceptualLoss().to(device).eval()
+    # We currently disable this as it is not used in the optimization.
+    perception_loss_module = None  # MultiLayerVGGPerceptualLoss().to(device).eval()
 
     # Create an optimizer instance
     optimizer = FilamentOptimizer(
@@ -300,7 +300,7 @@ def main():
     print("Starting optimization...")
     tbar = tqdm(range(args.iterations))
     for i in tbar:
-        loss_val = optimizer.step(record_best=i % 3 == 0)
+        loss_val = optimizer.step(record_best=i % 10 == 0)
 
         optimizer.visualize(interval=25)
         optimizer.log_to_tensorboard(interval=100)