autoforge 1.9.2__tar.gz → 1.9.2.2__tar.gz

{autoforge-1.9.2 → autoforge-1.9.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.9.2
+Version: 1.9.2.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

{autoforge-1.9.2 → autoforge-1.9.2.2}/pyproject.toml

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoForge
-Version: 1.9.2
+Version: 1.9.2.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

{autoforge-1.9.2 → autoforge-1.9.2.2}/src/AutoForge.egg-info/SOURCES.txt

@@ -10,6 +10,7 @@ src/autoforge/__init__.py
 src/autoforge/__main__.py
 src/autoforge/auto_forge.py
 src/autoforge/priority_mask_editor.py
+src/autoforge/Helper/AmpUtils.py
 src/autoforge/Helper/CAdamW.py
 src/autoforge/Helper/FilamentHelper.py
 src/autoforge/Helper/ImageHelper.py
@@ -32,6 +33,7 @@ src/autoforge/Misc/__init__.py
 src/autoforge/Misc/color_fit_all.py
 src/autoforge/Modules/Optimizer.py
 src/autoforge/Modules/__init__.py
+tests/test_amp_autocast_probe.py
 tests/test_cadamw.py
 tests/test_christofides_height_map.py
 tests/test_christofides_heightmap.py

autoforge-1.9.2.2/src/autoforge/Helper/AmpUtils.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from contextlib import nullcontext, contextmanager
+from typing import Optional, Tuple, Dict
+
+import os
+import torch
+
+# Cache the selection per device key so CPU/CUDA/MPS don't interfere
+_SELECTED_MAP: Dict[str, Tuple[Optional[torch.dtype], str]] = {}
+
+
+def _device_key(device: torch.device) -> str:
+    if device.type == "cuda":
+        idx = device.index if device.index is not None else torch.cuda.current_device()
+        return f"cuda:{idx}"
+    return device.type
+
+
+def _probe_backward(device: torch.device, dtype: torch.dtype) -> bool:
+    """Try a tiny forward+backward under autocast for given dtype on the device.
+
+    Returns True if successful, False if a RuntimeError occurs (e.g., oneDNN bf16/f16 backward not supported).
+    """
+    try:
+        if device.type == "mps":
+            return False  # no stable autocast on MPS yet
+        x = torch.randn(8, 8, device=device, dtype=torch.float32, requires_grad=False)
+        w = torch.randn(8, 8, device=device, dtype=torch.float32, requires_grad=True)
+        target = torch.randn(8, 8, device=device, dtype=torch.float32)
+        with torch.autocast(device_type=device.type, dtype=dtype):
+            y = x @ w
+            loss = torch.nn.functional.mse_loss(y, target)
+        loss.backward()
+        # Step with a dummy optimizer to ensure gradients are usable
+        opt = torch.optim.SGD([w], lr=1e-3)
+        opt.step()
+        return True
+    except Exception:
+        return False
+
+
+def _select_dtype_for_device(device: torch.device) -> Tuple[Optional[torch.dtype], str]:
+    """Decide the best autocast dtype for the current device, probing backward support.
+
+    Returns (dtype or None, reason string).
+    None means: do not use autocast; stick to full float32.
+    """
+    # Allow override via env for debugging
+    mode = os.getenv("AUTOFORGE_AMP", default="auto").lower()
+
+    if mode in {"off", "fp32", "float32"}:
+        print("AUTOFORGE_AMP=off: disabling automatic mixed precision")
+        return None, "AUTOFORGE_AMP=off"
+    if mode in {"bf16", "bfloat16"}:
+        # Only honor if probe succeeds
+        if _probe_backward(device, torch.bfloat16):
+            return torch.bfloat16, "forced bf16"
+        print("AUTOFORGE_AMP=bf16: disabling bf16 as probe failed")
+        return None, "forced bf16 but probe failed"
+    if mode in {"fp16", "float16", "half"}:
+        if device.type == "cuda" and _probe_backward(device, torch.float16):
+            return torch.float16, "forced f16"
+        print("AUTOFORGE_AMP=fp16: disabling fp16 as probe failed or non-CUDA")
+        return None, "forced f16 but probe failed or non-CUDA"
+
+    # Automatic selection
+    if device.type == "cuda":
+        # Prefer bf16 if it actually works end-to-end
+        if torch.cuda.is_bf16_supported() and _probe_backward(device, torch.bfloat16):
+            print("Using bf16 autocast on CUDA device")
+            return torch.bfloat16, "cuda bf16"
+        # Next try fp16 if the GPU supports it
+        major, minor = torch.cuda.get_device_capability(device)
+        if (major * 10 + minor) >= 53 and _probe_backward(device, torch.float16):
+            print("Using fp16 autocast on CUDA device")
+            return torch.float16, "cuda f16"
+        # Fall back to fp32; enable TF32 to speed up matmuls if available
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        print("Falling back to fp32 autocast on CUDA device with TF32 enabled")
+        return None, "cuda fp32 (tf32 on)"
+
+    if device.type == "cpu":
+        # Some CPUs (AVX512) can do bf16 forward/backward with oneDNN, but many cannot.
+        # Probe and use bf16 only if backward succeeds; otherwise stay fp32.
+        if _probe_backward(device, torch.bfloat16):
+            print("Using bf16 autocast on CPU device")
+            return torch.bfloat16, "cpu bf16"
+        # No reliable fp16 on CPU for backward; use fp32
+        print("Falling back to fp32 autocast on CPU device")
+        return None, "cpu fp32"
+
+    if device.type == "mps":
+        print("MPS device detected: autocast not reliably supported, using fp32")
+        return None, "mps fp32"
+
+    # Default: no autocast
+    print(f"Unknown device type '{device.type}': disabling autocast, using fp32")
+    return None, f"{device.type} fp32"
+
+
+def get_selected_autocast(device: torch.device) -> Tuple[Optional[torch.dtype], str]:
+    key = _device_key(device)
+    if key not in _SELECTED_MAP:
+        _SELECTED_MAP[key] = _select_dtype_for_device(device)
+    return _SELECTED_MAP[key]
+
+
+@contextmanager
+def safe_autocast(device: torch.device):
+    """Context manager that uses autocast only if the runtime probe succeeded.
+
+    Falls back to a no-op context on unsupported platforms to avoid runtime errors like:
+    RuntimeError: DNNL does not support bf16/f16 backward on the platform with avx2_vnni_2.
+    """
+    dtype, _reason = get_selected_autocast(device)
+    if dtype is None:
+        with nullcontext():
+            yield
+    else:
+        with torch.autocast(device_type=device.type, dtype=dtype):
+            yield

{autoforge-1.9.2 → autoforge-1.9.2.2}/src/autoforge/Helper/OptimizerHelper.py

@@ -2,6 +2,7 @@ from contextlib import contextmanager
 
 import torch
 import torch.nn.functional as F
+from autoforge.Helper.AmpUtils import get_selected_autocast
 
 
 @torch.jit.script
@@ -343,22 +344,25 @@ class PrecisionManager:
         self.autocast_dtype = None
         self.enabled = False
 
-        if device.type == "cuda":
-            if torch.cuda.is_bf16_supported():
-                print("Using bfloat16 autocast for CUDA.")
-                self.autocast_dtype = torch.bfloat16  # fastest if available
-                self.scaler = torch.cuda.amp.GradScaler()
-                self.enabled = True
-            elif _has_fp16(device):
-                print("Using float16 autocast for CUDA.")
-                self.autocast_dtype = torch.float16
-                self.scaler = torch.cuda.amp.GradScaler()
-                self.enabled = True
-            else:
-                print("Using float32 autocast for CUDA.")
-                # TF32 fallback for very old GPUs
+        # Decide dtype once using the shared runtime probe
+        dtype, _reason = get_selected_autocast(device)
+        self.autocast_dtype = dtype
+        # Enable only when a non-None dtype is selected and device supports native CUDA autocast
+        self.enabled = dtype is not None and device.type == "cuda"
+
+        # Use GradScaler only for CUDA float16; bf16 does not need scaling
+        if self.enabled and dtype == torch.float16:
+            self.scaler = torch.cuda.amp.GradScaler()
+        else:
+            self.scaler = None
+
+        # Optional: If CUDA but no AMP selected, allow TF32 for speed on Ampere+
+        if device.type == "cuda" and dtype is None:
+            try:
                 torch.backends.cuda.matmul.allow_tf32 = True
                 torch.backends.cudnn.allow_tf32 = True
+            except Exception:
+                pass
 
     @contextmanager
     def autocast(self):

{autoforge-1.9.2 → autoforge-1.9.2.2}/src/autoforge/auto_forge.py

@@ -191,7 +191,6 @@ def parse_args() -> argparse.Namespace:
 
     parser.add_argument(
         "--perform_pruning",
-        type=bool,
         default=True,
         help="Perform pruning after optimization",
         action=argparse.BooleanOptionalAction,
@@ -199,7 +198,6 @@ def parse_args() -> argparse.Namespace:
 
     parser.add_argument(
         "--fast_pruning",
-        type=bool,
         default=True,
         help="Use fast pruning method",
         action=argparse.BooleanOptionalAction,
@@ -289,7 +287,6 @@ def parse_args() -> argparse.Namespace:
 
     parser.add_argument(
         "--flatforge",
-        type=bool,
         default=False,
         help="Enable FlatForge mode to generate separate STL files for each color",
         action=argparse.BooleanOptionalAction,
@@ -461,7 +458,7 @@ def _prepare_background_and_materials(
     Returns:
         (bgr_tuple_uint8, background_tensor, material_colors_tensor, material_TDs_tensor)
     """
-    bgr_tuple = hex_to_rgb(args.background_color)
+    bgr_tuple = tuple(hex_to_rgb(args.background_color))
     background = torch.tensor(bgr_tuple, dtype=torch.float32, device=device)
     material_colors = torch.tensor(
         material_colors_np, dtype=torch.float32, device=device

autoforge-1.9.2.2/tests/test_amp_autocast_probe.py

@@ -0,0 +1,52 @@
+import os
+import torch
+import pytest
+
+from autoforge.Helper.AmpUtils import safe_autocast, get_selected_autocast
+
+
+def _device_from_env():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def test_safe_autocast_backward_does_not_crash():
+    # Ensure automatic mode
+    os.environ.pop("AUTOFORGE_AMP", None)
+    device = _device_from_env()
+
+    x = torch.randn(16, 16, device=device, dtype=torch.float32)
+    w = torch.randn(16, 16, device=device, dtype=torch.float32, requires_grad=True)
+    target = torch.randn(16, 16, device=device, dtype=torch.float32)
+
+    dtype, _reason = get_selected_autocast(device)
+    # Should not raise regardless of dtype selection
+    with safe_autocast(device):
+        y = x @ w
+        loss = torch.nn.functional.mse_loss(y, target)
+    loss.backward()
+
+    # If CUDA fp16 was selected, ensure grads exist and are finite
+    if device.type == "cuda" and dtype == torch.float16:
+        assert w.grad is not None
+        assert torch.isfinite(w.grad).all()
+
+
+@pytest.mark.parametrize("force", ["off", "bf16", "fp16"])  # exercise overrides
+def test_env_override_modes(monkeypatch, force):
+    monkeypatch.setenv("AUTOFORGE_AMP", force)
+    device = _device_from_env()
+
+    x = torch.randn(8, 8, device=device, dtype=torch.float32)
+    w = torch.randn(8, 8, device=device, dtype=torch.float32, requires_grad=True)
+    target = torch.randn(8, 8, device=device, dtype=torch.float32)
+
+    with safe_autocast(device):
+        y = x @ w
+        loss = torch.nn.functional.mse_loss(y, target)
+    loss.backward()
+    assert w.grad is not None
+