autoforge 1.9.6__tar.gz → 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (128) hide show
  1. {autoforge-1.9.6/src/AutoForge.egg-info → autoforge-2.0.0}/PKG-INFO +58 -9
  2. autoforge-1.9.6/PKG-INFO → autoforge-2.0.0/README.md +50 -38
  3. {autoforge-1.9.6 → autoforge-2.0.0}/pyproject.toml +14 -2
  4. autoforge-1.9.6/README.md → autoforge-2.0.0/src/AutoForge.egg-info/PKG-INFO +87 -7
  5. {autoforge-1.9.6 → autoforge-2.0.0}/src/AutoForge.egg-info/SOURCES.txt +18 -9
  6. {autoforge-1.9.6 → autoforge-2.0.0}/src/AutoForge.egg-info/requires.txt +7 -0
  7. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/AmpUtils.py +58 -19
  8. autoforge-2.0.0/src/autoforge/Helper/DeviceUtils.py +259 -0
  9. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/DepthEstimateHeightMap.py +21 -4
  10. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/FastTSPHeightMap.py +42 -24
  11. autoforge-2.0.0/src/autoforge/Helper/OptimizerHelper.py +822 -0
  12. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/OtherHelper.py +26 -7
  13. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/PruningHelper.py +88 -53
  14. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Loss/LossFunctions.py +17 -5
  15. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Modules/Optimizer.py +651 -86
  16. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/auto_forge.py +58 -7
  17. autoforge-2.0.0/src/autoforge/webui/api/filaments.py +231 -0
  18. autoforge-2.0.0/src/autoforge/webui/api/images.py +49 -0
  19. autoforge-2.0.0/src/autoforge/webui/api/init.py +359 -0
  20. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/api/optimization.py +147 -39
  21. autoforge-2.0.0/src/autoforge/webui/api/outputs.py +212 -0
  22. autoforge-2.0.0/src/autoforge/webui/api/preview.py +157 -0
  23. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/api/project.py +6 -1
  24. autoforge-2.0.0/src/autoforge/webui/api/pruning.py +415 -0
  25. autoforge-2.0.0/src/autoforge/webui/api/sliders.py +107 -0
  26. autoforge-2.0.0/src/autoforge/webui/api/state.py +63 -0
  27. autoforge-2.0.0/src/autoforge/webui/api/system.py +189 -0
  28. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/api/ws.py +28 -6
  29. autoforge-2.0.0/src/autoforge/webui/config.py +74 -0
  30. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/helpers/colored_mesh.py +141 -21
  31. autoforge-2.0.0/src/autoforge/webui/helpers/gpu_memory.py +80 -0
  32. autoforge-2.0.0/src/autoforge/webui/helpers/mesh_persist.py +78 -0
  33. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/helpers/pipeline_runner.py +237 -54
  34. autoforge-2.0.0/src/autoforge/webui/helpers/slider_render.py +341 -0
  35. autoforge-2.0.0/src/autoforge/webui/helpers/sliders.py +233 -0
  36. autoforge-2.0.0/src/autoforge/webui/helpers/telemetry.py +113 -0
  37. autoforge-2.0.0/src/autoforge/webui/models.py +219 -0
  38. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/server.py +14 -4
  39. autoforge-2.0.0/src/autoforge/webui/services/filament_service.py +358 -0
  40. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/services/image_service.py +14 -2
  41. autoforge-2.0.0/src/autoforge/webui/services/optimization_service.py +440 -0
  42. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/services/project_service.py +39 -1
  43. autoforge-2.0.0/tests/test_amp_utils.py +94 -0
  44. autoforge-2.0.0/tests/test_auto_forge_cli.py +60 -0
  45. autoforge-2.0.0/tests/test_bug_regressions.py +413 -0
  46. autoforge-2.0.0/tests/test_cadamw.py +148 -0
  47. autoforge-2.0.0/tests/test_christofides_heightmap.py +183 -0
  48. autoforge-2.0.0/tests/test_composite_image.py +269 -0
  49. {autoforge-1.9.6 → autoforge-2.0.0}/tests/test_depth_adjusted_height_map.py +1 -1
  50. autoforge-2.0.0/tests/test_device_utils.py +287 -0
  51. autoforge-2.0.0/tests/test_fasttsp_heightmap.py +149 -0
  52. autoforge-2.0.0/tests/test_filament_helper.py +186 -0
  53. autoforge-2.0.0/tests/test_image_helper.py +118 -0
  54. autoforge-2.0.0/tests/test_loss_functions.py +166 -0
  55. autoforge-2.0.0/tests/test_optimizer_training.py +361 -0
  56. autoforge-2.0.0/tests/test_other_helper.py +81 -0
  57. autoforge-2.0.0/tests/test_output_helper.py +160 -0
  58. autoforge-2.0.0/tests/test_perception_loss.py +65 -0
  59. autoforge-2.0.0/tests/test_pipeline_smoke.py +209 -0
  60. autoforge-2.0.0/tests/test_prune_repeatability.py +375 -0
  61. autoforge-2.0.0/tests/test_pruning_helper.py +102 -0
  62. autoforge-2.0.0/tests/test_webui_backend.py +505 -0
  63. autoforge-2.0.0/tests/test_webui_regressions.py +351 -0
  64. autoforge-1.9.6/src/autoforge/Helper/OptimizerHelper.py +0 -441
  65. autoforge-1.9.6/src/autoforge/webui/api/filaments.py +0 -120
  66. autoforge-1.9.6/src/autoforge/webui/api/images.py +0 -24
  67. autoforge-1.9.6/src/autoforge/webui/api/init.py +0 -37
  68. autoforge-1.9.6/src/autoforge/webui/api/outputs.py +0 -50
  69. autoforge-1.9.6/src/autoforge/webui/api/preview.py +0 -59
  70. autoforge-1.9.6/src/autoforge/webui/api/pruning.py +0 -154
  71. autoforge-1.9.6/src/autoforge/webui/api/sliders.py +0 -30
  72. autoforge-1.9.6/src/autoforge/webui/api/state.py +0 -38
  73. autoforge-1.9.6/src/autoforge/webui/api/system.py +0 -36
  74. autoforge-1.9.6/src/autoforge/webui/config.py +0 -27
  75. autoforge-1.9.6/src/autoforge/webui/helpers/slider_render.py +0 -202
  76. autoforge-1.9.6/src/autoforge/webui/helpers/sliders.py +0 -104
  77. autoforge-1.9.6/src/autoforge/webui/models.py +0 -116
  78. autoforge-1.9.6/src/autoforge/webui/services/filament_service.py +0 -219
  79. autoforge-1.9.6/src/autoforge/webui/services/optimization_service.py +0 -193
  80. autoforge-1.9.6/tests/test_amp_autocast_probe.py +0 -52
  81. autoforge-1.9.6/tests/test_cadamw.py +0 -50
  82. autoforge-1.9.6/tests/test_christofides_height_map.py +0 -63
  83. autoforge-1.9.6/tests/test_christofides_heightmap.py +0 -70
  84. autoforge-1.9.6/tests/test_filament_and_output.py +0 -105
  85. autoforge-1.9.6/tests/test_filament_helper.py +0 -75
  86. autoforge-1.9.6/tests/test_heightmap_init.py +0 -54
  87. autoforge-1.9.6/tests/test_image_helper.py +0 -42
  88. autoforge-1.9.6/tests/test_loss_functions.py +0 -33
  89. autoforge-1.9.6/tests/test_optimizer_helper.py +0 -65
  90. autoforge-1.9.6/tests/test_optimizer_smoke.py +0 -59
  91. autoforge-1.9.6/tests/test_output_helper.py +0 -105
  92. autoforge-1.9.6/tests/test_output_helper_extended.py +0 -88
  93. autoforge-1.9.6/tests/test_perception_and_pruning.py +0 -103
  94. autoforge-1.9.6/tests/test_pruning_helper.py +0 -89
  95. {autoforge-1.9.6 → autoforge-2.0.0}/setup.cfg +0 -0
  96. {autoforge-1.9.6 → autoforge-2.0.0}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  97. {autoforge-1.9.6 → autoforge-2.0.0}/src/AutoForge.egg-info/entry_points.txt +0 -0
  98. {autoforge-1.9.6 → autoforge-2.0.0}/src/AutoForge.egg-info/top_level.txt +0 -0
  99. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/CAdamW.py +0 -0
  100. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/FilamentHelper.py +0 -0
  101. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/ChristofidesHeightMap.py +0 -0
  102. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/__init__.py +0 -0
  103. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/ImageHelper.py +0 -0
  104. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/OutputHelper.py +0 -0
  105. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/__init__.py +0 -0
  106. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Helper/filamentcolors_library.py +0 -0
  107. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Init_comparer/__init__.py +0 -0
  108. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Init_comparer/auto_forge_init_check.py +0 -0
  109. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Init_comparer/plot_results.py +0 -0
  110. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  111. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Loss/__init__.py +0 -0
  112. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Misc/__init__.py +0 -0
  113. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Misc/color_fit_all.py +0 -0
  114. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/Modules/__init__.py +0 -0
  115. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/__init__.py +0 -0
  116. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/__main__.py +0 -0
  117. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/priority_mask_editor.py +0 -0
  118. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/__init__.py +0 -0
  119. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/api/__init__.py +0 -0
  120. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/api/settings.py +0 -0
  121. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/helpers/__init__.py +0 -0
  122. {autoforge-1.9.6 → autoforge-2.0.0}/src/autoforge/webui/services/__init__.py +0 -0
  123. {autoforge-1.9.6 → autoforge-2.0.0}/tests/test_depth_estimate_height_map.py +0 -0
  124. {autoforge-1.9.6 → autoforge-2.0.0}/tests/test_end_to_end_optimizer.py +0 -0
  125. {autoforge-1.9.6 → autoforge-2.0.0}/tests/test_flatforge.py +0 -0
  126. {autoforge-1.9.6 → autoforge-2.0.0}/tests/test_optimizer_module.py +0 -0
  127. {autoforge-1.9.6 → autoforge-2.0.0}/tests/test_property_based.py +0 -0
  128. {autoforge-1.9.6 → autoforge-2.0.0}/tests/test_remove_height_spikes.py +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: AutoForge
3
- Version: 1.9.6
4
- Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
3
+ Version: 2.0.0
4
+ Summary: AutoForge - 3D color images made easy
5
5
  Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
6
  License: CC BY-NC-SA 4.0
7
7
  Requires-Python: >=3.10
@@ -28,17 +28,22 @@ Requires-Dist: python-multipart>=0.0.12
28
28
  Requires-Dist: aiofiles>=24.1.0
29
29
  Requires-Dist: websockets>=15.0
30
30
  Requires-Dist: pydantic-settings>=2.0.0
31
+ Requires-Dist: posthog>=3.7.0
32
+ Provides-Extra: test
33
+ Requires-Dist: pytest>=7.4; extra == "test"
34
+ Requires-Dist: pytest-cov; extra == "test"
35
+ Requires-Dist: hypothesis; extra == "test"
36
+ Requires-Dist: httpx>=0.27; extra == "test"
31
37
 
32
38
  # AutoForge
33
39
 
34
40
  AutoForge is a Python tool for generating 3D printed layered models from an input image. Using a learned optimization strategy with a Gumbel softmax formulation, AutoForge assigns materials per layer and produces both a discretized composite image and a 3D-printable STL file. It also generates swap instructions to guide the printer through material changes during a multi-material print.
35
41
 
36
- **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).
42
+ **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.
43
+
44
+ ![2.0_webui.png](images/webui/2.0_webui.png)
37
45
 
38
46
 
39
- ## You can now run Autoforge for free in your browser thanks to [Huggingface space support](https://huggingface.co/spaces/hvoss-techfak/Autoforge).
40
- This includes the option to run it locally if you have a powerful pc and don't want to limit yourself to the Huggingface computing limits. \
41
- For this simply go to the [Huggingface](https://huggingface.co/spaces/hvoss-techfak/Autoforge) space and pull the docker container for this project (upper right corner -> three dots -> "run locally")
42
47
 
43
48
  ## Example
44
49
  All examples use only the 27 BambuLab Basic PLA filaments, currently available in Hueforge 0.9.0, the background color is set to black.
@@ -72,9 +77,41 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
72
77
  - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
73
78
  - **Hueforge export**: Outputs a project file that can be opened with hueforge.
74
79
 
75
- ## Manual Installation
80
+ ## Web UI: One-Click Install & Run
81
+
82
+
83
+ The easiest way to use AutoForge is the web UI, a local app with drag-and-drop image upload, a filament library, live sliders for adjusting colors after optimization, and pruning. No command-line arguments needed.
84
+
85
+ 1. **Clone this repository** (or download and extract the ZIP from the green "Code" button on GitHub).
86
+ 2. **Install**, from the project folder:
87
+ - Linux/macOS: `./install.sh`
88
+ - Windows: double-click `install.bat` (or run it from a terminal)
89
+
90
+ This installs [`uv`](https://docs.astral.sh/uv/) (a fast Python package manager) if you don't already have it, installs all Python dependencies, and builds the web UI. You'll need [Node.js](https://nodejs.org/) installed for that last step, the installer will tell you if it's missing.
91
+ 3. **Run**:
92
+ - Linux/macOS: `./run_webui.sh`
93
+ - Windows: double-click `run_webui.bat`
94
+
95
+ This starts the server and opens the web UI in your browser automatically (usually at `http://localhost:8000`).
96
+
97
+ To make some parts of the picture come out closer than the rest (a face, the eyes, lettering), click **Focus** in the image panel and paint over them; a slider sets how much more they count (2× to 100×, default 10×). After a run, the **Differences** view shows where the print strays furthest from the picture.
98
+
99
+ If HueForge is installed on the same computer, the web UI finds your HueForge personal filament library and offers to import it once, on its first start.
76
100
 
77
- To install AutoForge, simply install the current version from PyPI:
101
+ The web UI sends anonymous usage telemetry to the project via [PostHog](https://posthog.com/) by default, to notify me of problems and any bugs. This includes crash reports (unhandled errors from both the browser frontend and the backend server, with the error type, message, and stack trace) so bugs can get fixed faster. No image data, filament data, or personal information is sent. To disable it, pass `--no-telemetry` (e.g. `./run_webui.sh --no-telemetry` / `run_webui.bat --no-telemetry`), or set `AUTOFORGE_WEBUI_TELEMETRY_ENABLED=false` permanently in your environment.
102
+
103
+ 4. **Update to the latest release** whenever you want, from the project folder:
104
+ - Linux/macOS: `./update.sh`
105
+ - Windows: double-click `update.bat`
106
+
107
+ This checks GitHub for a newer release and, if one exists, pulls it and reinstalls dependencies for you. Add `--check` to only check without applying it (e.g. `./update.sh --check`).
108
+
109
+ If you have problems running the code on your GPU, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
110
+ CUDA, ROCm, and MPS (Apple Metal) are supported, but you need to install the correct version of pytorch for your system — `install.sh`/`install.bat` install whatever `uv sync` resolves by default, so swap in a GPU-specific PyTorch build afterwards if you need one. (`install.sh` handles older NVIDIA GPUs automatically, see below.)
111
+
112
+ ## Manual Installation (CLI only)
113
+
114
+ If you just want the command-line tool (no web UI), install the current version from PyPI:
78
115
  ```bash
79
116
  pip install -U autoforge
80
117
  ```
@@ -146,6 +183,8 @@ This will generate separate STL files for each color, allowing you to print face
146
183
  - `--nozzle_diameter` Diameter of the printer nozzle in millimeters (default: 0.4).
147
184
  **Note:** Details smaller than half this value will be ignored.
148
185
  - `--early_stopping` Number of steps without improvement before stopping (default: 10000).
186
+ - `--priority_mask` *(Optional)* Path to a greyscale image the size of the input that marks the parts that matter most: white areas are matched more closely, black areas still count, just less.
187
+ - `--priority_mask_strength` How many times more a white pixel of `--priority_mask` counts than a black one (default: 10; must be at least 1). Raise it if the marked areas still come out off, lower it for a gentler nudge.
149
188
 
150
189
  - `--flatforge` Enable FlatForge mode to generate separate STL files for each color (default: False).
151
190
  **Note:** FlatForge creates flat prints where each color is its own STL file, allowing face-down printing for smooth, resin-like finishes. Requires a multi-material printer (AMS, MMU, or tool changer).
@@ -163,7 +202,8 @@ This will generate separate STL files for each color, allowing you to print face
163
202
  - `--pruning_max_swaps` Max number of swaps allowed after pruning (default: 100).
164
203
  - `--pruning_max_layer` Max number of layers allowed after pruning (default: 75).
165
204
  - `--random_seed` Random seed for reproducibility (default: 0 (disabled)).
166
- - `--mps` Flag to use the Metal Performance Shaders (MPS) backend if available.
205
+ - `--device` Torch device to run on, e.g. `cuda`, `cuda:1`, `mps`, `cpu`. Defaults to auto-detection: CUDA/ROCm first, then Apple Metal (MPS), then CPU. Can also be set with the `AUTOFORGE_DEVICE` environment variable.
206
+ - `--mps` *Deprecated* — Apple Metal is now detected automatically, so this flag is no longer needed. It still works, and forces MPS on a machine that also exposes a CUDA GPU; prefer `--device mps`.
167
207
  - `--no-spike-removal` Disable spike removal for the final STL (not recommended).
168
208
 
169
209
  - `--tensorboard` Flag to enable TensorBoard logging.
@@ -192,6 +232,10 @@ After running, the following files will be created in your specified output fold
192
232
 
193
233
  *Note:* FlatForge generates multiple STL files that align perfectly when loaded together in your slicer, creating a solid rectangular print with each color as a separate object.
194
234
 
235
+ ## You can now run Autoforge for free in your browser thanks to [Huggingface space support](https://huggingface.co/spaces/hvoss-techfak/Autoforge).
236
+ This includes the option to run it locally if you have a powerful pc and don't want to limit yourself to the Huggingface computing limits. \
237
+ For this simply go to the [Huggingface](https://huggingface.co/spaces/hvoss-techfak/Autoforge) space and pull the docker container for this project (upper right corner -> three dots -> "run locally")
238
+
195
239
  ## Development
196
240
 
197
241
  To have a "nightly" version of the repository or have live updating changes during development please do the following:
@@ -221,6 +265,11 @@ pip install torch torchvision torchaudio --index-url https://download.pytorch.or
221
265
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
222
266
  The software is provided as-is and comes with no warranty or guarantee of support.
223
267
 
268
+ The above license applies to the software itself. How you use the generated files and prints is entirely up to you. If you want to print them for your friends or family, that's great. If you want to sell them, that's fine by me too.
269
+
270
+ From a licensing standpoint, this means that the prints you create are entirely subject to the [MIT License](https://opensource.org/licenses/MIT), and you can do whatever you like with them. The only thing the MIT License does not give you is a warranty, and it also frees me from any liability with regard to your 3D prints. Otherwise, do whatever you like.
271
+
272
+ I would love to see what you have done with the software, so it would be great if you could send me a link to your work (even if it's just for selling). However, this is not necessary if you do not want to.
224
273
 
225
274
  ## Acknowledgements
226
275
 
@@ -1,44 +1,12 @@
1
- Metadata-Version: 2.4
2
- Name: AutoForge
3
- Version: 1.9.6
4
- Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
5
- Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
- License: CC BY-NC-SA 4.0
7
- Requires-Python: >=3.10
8
- Description-Content-Type: text/markdown
9
- Requires-Dist: numpy
10
- Requires-Dist: torch>=2.9.1
11
- Requires-Dist: torchvision>=0.21.0
12
- Requires-Dist: transformers>=4.48.1
13
- Requires-Dist: opencv-python>=4.2.0.32
14
- Requires-Dist: matplotlib>=3.3.0
15
- Requires-Dist: configargparse>=1.2.3
16
- Requires-Dist: tqdm>=4.50.0
17
- Requires-Dist: pandas>=1.1.0
18
- Requires-Dist: scikit-image
19
- Requires-Dist: scikit-learn
20
- Requires-Dist: tensorboard
21
- Requires-Dist: pillow>=11.1.0
22
- Requires-Dist: shapely>=2.0.7
23
- Requires-Dist: trimesh>=4.6.4
24
- Requires-Dist: joblib>=1.4.2
25
- Requires-Dist: uvicorn[standard]>=0.44.0
26
- Requires-Dist: fastapi>=0.115.0
27
- Requires-Dist: python-multipart>=0.0.12
28
- Requires-Dist: aiofiles>=24.1.0
29
- Requires-Dist: websockets>=15.0
30
- Requires-Dist: pydantic-settings>=2.0.0
31
-
32
1
  # AutoForge
33
2
 
34
3
  AutoForge is a Python tool for generating 3D printed layered models from an input image. Using a learned optimization strategy with a Gumbel softmax formulation, AutoForge assigns materials per layer and produces both a discretized composite image and a 3D-printable STL file. It also generates swap instructions to guide the printer through material changes during a multi-material print.
35
4
 
36
- **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).
5
+ **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.
6
+
7
+ ![2.0_webui.png](images/webui/2.0_webui.png)
37
8
 
38
9
 
39
- ## You can now run Autoforge for free in your browser thanks to [Huggingface space support](https://huggingface.co/spaces/hvoss-techfak/Autoforge).
40
- This includes the option to run it locally if you have a powerful pc and don't want to limit yourself to the Huggingface computing limits. \
41
- For this simply go to the [Huggingface](https://huggingface.co/spaces/hvoss-techfak/Autoforge) space and pull the docker container for this project (upper right corner -> three dots -> "run locally")
42
10
 
43
11
  ## Example
44
12
  All examples use only the 27 BambuLab Basic PLA filaments, currently available in Hueforge 0.9.0, the background color is set to black.
@@ -72,9 +40,41 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
72
40
  - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
73
41
  - **Hueforge export**: Outputs a project file that can be opened with hueforge.
74
42
 
75
- ## Manual Installation
43
+ ## Web UI: One-Click Install & Run
44
+
45
+
46
+ The easiest way to use AutoForge is the web UI, a local app with drag-and-drop image upload, a filament library, live sliders for adjusting colors after optimization, and pruning. No command-line arguments needed.
47
+
48
+ 1. **Clone this repository** (or download and extract the ZIP from the green "Code" button on GitHub).
49
+ 2. **Install**, from the project folder:
50
+ - Linux/macOS: `./install.sh`
51
+ - Windows: double-click `install.bat` (or run it from a terminal)
52
+
53
+ This installs [`uv`](https://docs.astral.sh/uv/) (a fast Python package manager) if you don't already have it, installs all Python dependencies, and builds the web UI. You'll need [Node.js](https://nodejs.org/) installed for that last step, the installer will tell you if it's missing.
54
+ 3. **Run**:
55
+ - Linux/macOS: `./run_webui.sh`
56
+ - Windows: double-click `run_webui.bat`
57
+
58
+ This starts the server and opens the web UI in your browser automatically (usually at `http://localhost:8000`).
59
+
60
+ To make some parts of the picture come out closer than the rest (a face, the eyes, lettering), click **Focus** in the image panel and paint over them; a slider sets how much more they count (2× to 100×, default 10×). After a run, the **Differences** view shows where the print strays furthest from the picture.
61
+
62
+ If HueForge is installed on the same computer, the web UI finds your HueForge personal filament library and offers to import it once, on its first start.
76
63
 
77
- To install AutoForge, simply install the current version from PyPI:
64
+ The web UI sends anonymous usage telemetry to the project via [PostHog](https://posthog.com/) by default, to notify me of problems and any bugs. This includes crash reports (unhandled errors from both the browser frontend and the backend server, with the error type, message, and stack trace) so bugs can get fixed faster. No image data, filament data, or personal information is sent. To disable it, pass `--no-telemetry` (e.g. `./run_webui.sh --no-telemetry` / `run_webui.bat --no-telemetry`), or set `AUTOFORGE_WEBUI_TELEMETRY_ENABLED=false` permanently in your environment.
65
+
66
+ 4. **Update to the latest release** whenever you want, from the project folder:
67
+ - Linux/macOS: `./update.sh`
68
+ - Windows: double-click `update.bat`
69
+
70
+ This checks GitHub for a newer release and, if one exists, pulls it and reinstalls dependencies for you. Add `--check` to only check without applying it (e.g. `./update.sh --check`).
71
+
72
+ If you have problems running the code on your GPU, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
73
+ CUDA, ROCm, and MPS (Apple Metal) are supported, but you need to install the correct version of pytorch for your system — `install.sh`/`install.bat` install whatever `uv sync` resolves by default, so swap in a GPU-specific PyTorch build afterwards if you need one. (`install.sh` handles older NVIDIA GPUs automatically, see below.)
74
+
75
+ ## Manual Installation (CLI only)
76
+
77
+ If you just want the command-line tool (no web UI), install the current version from PyPI:
78
78
  ```bash
79
79
  pip install -U autoforge
80
80
  ```
@@ -146,6 +146,8 @@ This will generate separate STL files for each color, allowing you to print face
146
146
  - `--nozzle_diameter` Diameter of the printer nozzle in millimeters (default: 0.4).
147
147
  **Note:** Details smaller than half this value will be ignored.
148
148
  - `--early_stopping` Number of steps without improvement before stopping (default: 10000).
149
+ - `--priority_mask` *(Optional)* Path to a greyscale image the size of the input that marks the parts that matter most: white areas are matched more closely, black areas still count, just less.
150
+ - `--priority_mask_strength` How many times more a white pixel of `--priority_mask` counts than a black one (default: 10; must be at least 1). Raise it if the marked areas still come out off, lower it for a gentler nudge.
149
151
 
150
152
  - `--flatforge` Enable FlatForge mode to generate separate STL files for each color (default: False).
151
153
  **Note:** FlatForge creates flat prints where each color is its own STL file, allowing face-down printing for smooth, resin-like finishes. Requires a multi-material printer (AMS, MMU, or tool changer).
@@ -163,7 +165,8 @@ This will generate separate STL files for each color, allowing you to print face
163
165
  - `--pruning_max_swaps` Max number of swaps allowed after pruning (default: 100).
164
166
  - `--pruning_max_layer` Max number of layers allowed after pruning (default: 75).
165
167
  - `--random_seed` Random seed for reproducibility (default: 0 (disabled)).
166
- - `--mps` Flag to use the Metal Performance Shaders (MPS) backend if available.
168
+ - `--device` Torch device to run on, e.g. `cuda`, `cuda:1`, `mps`, `cpu`. Defaults to auto-detection: CUDA/ROCm first, then Apple Metal (MPS), then CPU. Can also be set with the `AUTOFORGE_DEVICE` environment variable.
169
+ - `--mps` *Deprecated* — Apple Metal is now detected automatically, so this flag is no longer needed. It still works, and forces MPS on a machine that also exposes a CUDA GPU; prefer `--device mps`.
167
170
  - `--no-spike-removal` Disable spike removal for the final STL (not recommended).
168
171
 
169
172
  - `--tensorboard` Flag to enable TensorBoard logging.
@@ -192,6 +195,10 @@ After running, the following files will be created in your specified output fold
192
195
 
193
196
  *Note:* FlatForge generates multiple STL files that align perfectly when loaded together in your slicer, creating a solid rectangular print with each color as a separate object.
194
197
 
198
+ ## You can now run Autoforge for free in your browser thanks to [Huggingface space support](https://huggingface.co/spaces/hvoss-techfak/Autoforge).
199
+ This includes the option to run it locally if you have a powerful pc and don't want to limit yourself to the Huggingface computing limits. \
200
+ For this simply go to the [Huggingface](https://huggingface.co/spaces/hvoss-techfak/Autoforge) space and pull the docker container for this project (upper right corner -> three dots -> "run locally")
201
+
195
202
  ## Development
196
203
 
197
204
  To have a "nightly" version of the repository or have live updating changes during development please do the following:
@@ -221,6 +228,11 @@ pip install torch torchvision torchaudio --index-url https://download.pytorch.or
221
228
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
222
229
  The software is provided as-is and comes with no warranty or guarantee of support.
223
230
 
231
+ The above license applies to the software itself. How you use the generated files and prints is entirely up to you. If you want to print them for your friends or family, that's great. If you want to sell them, that's fine by me too.
232
+
233
+ From a licensing standpoint, this means that the prints you create are entirely subject to the [MIT License](https://opensource.org/licenses/MIT), and you can do whatever you like with them. The only thing the MIT License does not give you is a warranty, and it also frees me from any liability with regard to your 3D prints. Otherwise, do whatever you like.
234
+
235
+ I would love to see what you have done with the software, so it would be great if you could send me a link to your work (even if it's just for selling). However, this is not necessary if you do not want to.
224
236
 
225
237
  ## Acknowledgements
226
238
 
@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="1.9.6"
8
- description = "AutoForge is a Python tool for generating 3D printed layered models from an input image."
7
+ version="2.0.0"
8
+ description = "AutoForge - 3D color images made easy"
9
9
  authors = [{ name = "Hendric Voss", email = "hvoss@techfak.uni-bielefeld.de" }]
10
10
  readme = "README.md"
11
11
  requires-python = ">=3.10"
@@ -33,6 +33,18 @@ dependencies = [
33
33
  "aiofiles>=24.1.0",
34
34
  "websockets>=15.0",
35
35
  "pydantic-settings>=2.0.0",
36
+ "posthog>=3.7.0",
37
+ ]
38
+
39
+ [project.optional-dependencies]
40
+ # httpx isn't a runtime dependency of the app itself, but fastapi's
41
+ # TestClient needs it installed, and tests/webui/test_review_2026_09_23.py
42
+ # uses httpx.AsyncClient/ASGITransport directly to test concurrent requests.
43
+ test = [
44
+ "pytest>=7.4",
45
+ "pytest-cov",
46
+ "hypothesis",
47
+ "httpx>=0.27",
36
48
  ]
37
49
 
38
50
  [tool.setuptools]
@@ -1,13 +1,49 @@
1
+ Metadata-Version: 2.4
2
+ Name: AutoForge
3
+ Version: 2.0.0
4
+ Summary: AutoForge - 3D color images made easy
5
+ Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
+ License: CC BY-NC-SA 4.0
7
+ Requires-Python: >=3.10
8
+ Description-Content-Type: text/markdown
9
+ Requires-Dist: numpy
10
+ Requires-Dist: torch>=2.9.1
11
+ Requires-Dist: torchvision>=0.21.0
12
+ Requires-Dist: transformers>=4.48.1
13
+ Requires-Dist: opencv-python>=4.2.0.32
14
+ Requires-Dist: matplotlib>=3.3.0
15
+ Requires-Dist: configargparse>=1.2.3
16
+ Requires-Dist: tqdm>=4.50.0
17
+ Requires-Dist: pandas>=1.1.0
18
+ Requires-Dist: scikit-image
19
+ Requires-Dist: scikit-learn
20
+ Requires-Dist: tensorboard
21
+ Requires-Dist: pillow>=11.1.0
22
+ Requires-Dist: shapely>=2.0.7
23
+ Requires-Dist: trimesh>=4.6.4
24
+ Requires-Dist: joblib>=1.4.2
25
+ Requires-Dist: uvicorn[standard]>=0.44.0
26
+ Requires-Dist: fastapi>=0.115.0
27
+ Requires-Dist: python-multipart>=0.0.12
28
+ Requires-Dist: aiofiles>=24.1.0
29
+ Requires-Dist: websockets>=15.0
30
+ Requires-Dist: pydantic-settings>=2.0.0
31
+ Requires-Dist: posthog>=3.7.0
32
+ Provides-Extra: test
33
+ Requires-Dist: pytest>=7.4; extra == "test"
34
+ Requires-Dist: pytest-cov; extra == "test"
35
+ Requires-Dist: hypothesis; extra == "test"
36
+ Requires-Dist: httpx>=0.27; extra == "test"
37
+
1
38
  # AutoForge
2
39
 
3
40
  AutoForge is a Python tool for generating 3D printed layered models from an input image. Using a learned optimization strategy with a Gumbel softmax formulation, AutoForge assigns materials per layer and produces both a discretized composite image and a 3D-printable STL file. It also generates swap instructions to guide the printer through material changes during a multi-material print.
4
41
 
5
- **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).
42
+ **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.
43
+
44
+ ![2.0_webui.png](images/webui/2.0_webui.png)
6
45
 
7
46
 
8
- ## You can now run Autoforge for free in your browser thanks to [Huggingface space support](https://huggingface.co/spaces/hvoss-techfak/Autoforge).
9
- This includes the option to run it locally if you have a powerful pc and don't want to limit yourself to the Huggingface computing limits. \
10
- For this simply go to the [Huggingface](https://huggingface.co/spaces/hvoss-techfak/Autoforge) space and pull the docker container for this project (upper right corner -> three dots -> "run locally")
11
47
 
12
48
  ## Example
13
49
  All examples use only the 27 BambuLab Basic PLA filaments, currently available in Hueforge 0.9.0, the background color is set to black.
@@ -41,9 +77,41 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
41
77
  - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
42
78
  - **Hueforge export**: Outputs a project file that can be opened with hueforge.
43
79
 
44
- ## Manual Installation
80
+ ## Web UI: One-Click Install & Run
81
+
82
+
83
+ The easiest way to use AutoForge is the web UI, a local app with drag-and-drop image upload, a filament library, live sliders for adjusting colors after optimization, and pruning. No command-line arguments needed.
84
+
85
+ 1. **Clone this repository** (or download and extract the ZIP from the green "Code" button on GitHub).
86
+ 2. **Install**, from the project folder:
87
+ - Linux/macOS: `./install.sh`
88
+ - Windows: double-click `install.bat` (or run it from a terminal)
89
+
90
+ This installs [`uv`](https://docs.astral.sh/uv/) (a fast Python package manager) if you don't already have it, installs all Python dependencies, and builds the web UI. You'll need [Node.js](https://nodejs.org/) installed for that last step, the installer will tell you if it's missing.
91
+ 3. **Run**:
92
+ - Linux/macOS: `./run_webui.sh`
93
+ - Windows: double-click `run_webui.bat`
94
+
95
+ This starts the server and opens the web UI in your browser automatically (usually at `http://localhost:8000`).
96
+
97
+ To make some parts of the picture come out closer than the rest (a face, the eyes, lettering), click **Focus** in the image panel and paint over them; a slider sets how much more they count (2× to 100×, default 10×). After a run, the **Differences** view shows where the print strays furthest from the picture.
98
+
99
+ If HueForge is installed on the same computer, the web UI finds your HueForge personal filament library and offers to import it once, on its first start.
45
100
 
46
- To install AutoForge, simply install the current version from PyPI:
101
+ The web UI sends anonymous usage telemetry to the project via [PostHog](https://posthog.com/) by default, to notify me of problems and any bugs. This includes crash reports (unhandled errors from both the browser frontend and the backend server, with the error type, message, and stack trace) so bugs can get fixed faster. No image data, filament data, or personal information is sent. To disable it, pass `--no-telemetry` (e.g. `./run_webui.sh --no-telemetry` / `run_webui.bat --no-telemetry`), or set `AUTOFORGE_WEBUI_TELEMETRY_ENABLED=false` permanently in your environment.
102
+
103
+ 4. **Update to the latest release** whenever you want, from the project folder:
104
+ - Linux/macOS: `./update.sh`
105
+ - Windows: double-click `update.bat`
106
+
107
+ This checks GitHub for a newer release and, if one exists, pulls it and reinstalls dependencies for you. Add `--check` to only check without applying it (e.g. `./update.sh --check`).
108
+
109
+ If you have problems running the code on your GPU, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
110
+ CUDA, ROCm, and MPS (Apple Metal) are supported, but you need to install the correct version of pytorch for your system — `install.sh`/`install.bat` install whatever `uv sync` resolves by default, so swap in a GPU-specific PyTorch build afterwards if you need one. (`install.sh` handles older NVIDIA GPUs automatically, see below.)
111
+
112
+ ## Manual Installation (CLI only)
113
+
114
+ If you just want the command-line tool (no web UI), install the current version from PyPI:
47
115
  ```bash
48
116
  pip install -U autoforge
49
117
  ```
@@ -115,6 +183,8 @@ This will generate separate STL files for each color, allowing you to print face
115
183
  - `--nozzle_diameter` Diameter of the printer nozzle in millimeters (default: 0.4).
116
184
  **Note:** Details smaller than half this value will be ignored.
117
185
  - `--early_stopping` Number of steps without improvement before stopping (default: 10000).
186
+ - `--priority_mask` *(Optional)* Path to a greyscale image the size of the input that marks the parts that matter most: white areas are matched more closely, black areas still count, just less.
187
+ - `--priority_mask_strength` How many times more a white pixel of `--priority_mask` counts than a black one (default: 10; must be at least 1). Raise it if the marked areas still come out off, lower it for a gentler nudge.
118
188
 
119
189
  - `--flatforge` Enable FlatForge mode to generate separate STL files for each color (default: False).
120
190
  **Note:** FlatForge creates flat prints where each color is its own STL file, allowing face-down printing for smooth, resin-like finishes. Requires a multi-material printer (AMS, MMU, or tool changer).
@@ -132,7 +202,8 @@ This will generate separate STL files for each color, allowing you to print face
132
202
  - `--pruning_max_swaps` Max number of swaps allowed after pruning (default: 100).
133
203
  - `--pruning_max_layer` Max number of layers allowed after pruning (default: 75).
134
204
  - `--random_seed` Random seed for reproducibility (default: 0 (disabled)).
135
- - `--mps` Flag to use the Metal Performance Shaders (MPS) backend if available.
205
+ - `--device` Torch device to run on, e.g. `cuda`, `cuda:1`, `mps`, `cpu`. Defaults to auto-detection: CUDA/ROCm first, then Apple Metal (MPS), then CPU. Can also be set with the `AUTOFORGE_DEVICE` environment variable.
206
+ - `--mps` *Deprecated* — Apple Metal is now detected automatically, so this flag is no longer needed. It still works, and forces MPS on a machine that also exposes a CUDA GPU; prefer `--device mps`.
136
207
  - `--no-spike-removal` Disable spike removal for the final STL (not recommended).
137
208
 
138
209
  - `--tensorboard` Flag to enable TensorBoard logging.
@@ -161,6 +232,10 @@ After running, the following files will be created in your specified output fold
161
232
 
162
233
  *Note:* FlatForge generates multiple STL files that align perfectly when loaded together in your slicer, creating a solid rectangular print with each color as a separate object.
163
234
 
235
+ ## You can now run Autoforge for free in your browser thanks to [Huggingface space support](https://huggingface.co/spaces/hvoss-techfak/Autoforge).
236
+ This includes the option to run it locally if you have a powerful pc and don't want to limit yourself to the Huggingface computing limits. \
237
+ For this simply go to the [Huggingface](https://huggingface.co/spaces/hvoss-techfak/Autoforge) space and pull the docker container for this project (upper right corner -> three dots -> "run locally")
238
+
164
239
  ## Development
165
240
 
166
241
  To have a "nightly" version of the repository or have live updating changes during development please do the following:
@@ -190,6 +265,11 @@ pip install torch torchvision torchaudio --index-url https://download.pytorch.or
190
265
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
191
266
  The software is provided as-is and comes with no warranty or guarantee of support.
192
267
 
268
+ The above license applies to the software itself. How you use the generated files and prints is entirely up to you. If you want to print them for your friends or family, that's great. If you want to sell them, that's fine by me too.
269
+
270
+ From a licensing standpoint, this means that the prints you create are entirely subject to the [MIT License](https://opensource.org/licenses/MIT), and you can do whatever you like with them. The only thing the MIT License does not give you is a warranty, and it also frees me from any liability with regard to your 3D prints. Otherwise, do whatever you like.
271
+
272
+ I would love to see what you have done with the software, so it would be great if you could send me a link to your work (even if it's just for selling). However, this is not necessary if you do not want to.
193
273
 
194
274
  ## Acknowledgements
195
275
 
@@ -12,6 +12,7 @@ src/autoforge/auto_forge.py
12
12
  src/autoforge/priority_mask_editor.py
13
13
  src/autoforge/Helper/AmpUtils.py
14
14
  src/autoforge/Helper/CAdamW.py
15
+ src/autoforge/Helper/DeviceUtils.py
15
16
  src/autoforge/Helper/FilamentHelper.py
16
17
  src/autoforge/Helper/ImageHelper.py
17
18
  src/autoforge/Helper/OptimizerHelper.py
@@ -54,33 +55,41 @@ src/autoforge/webui/api/system.py
54
55
  src/autoforge/webui/api/ws.py
55
56
  src/autoforge/webui/helpers/__init__.py
56
57
  src/autoforge/webui/helpers/colored_mesh.py
58
+ src/autoforge/webui/helpers/gpu_memory.py
59
+ src/autoforge/webui/helpers/mesh_persist.py
57
60
  src/autoforge/webui/helpers/pipeline_runner.py
58
61
  src/autoforge/webui/helpers/slider_render.py
59
62
  src/autoforge/webui/helpers/sliders.py
63
+ src/autoforge/webui/helpers/telemetry.py
60
64
  src/autoforge/webui/services/__init__.py
61
65
  src/autoforge/webui/services/filament_service.py
62
66
  src/autoforge/webui/services/image_service.py
63
67
  src/autoforge/webui/services/optimization_service.py
64
68
  src/autoforge/webui/services/project_service.py
65
- tests/test_amp_autocast_probe.py
69
+ tests/test_amp_utils.py
70
+ tests/test_auto_forge_cli.py
71
+ tests/test_bug_regressions.py
66
72
  tests/test_cadamw.py
67
- tests/test_christofides_height_map.py
68
73
  tests/test_christofides_heightmap.py
74
+ tests/test_composite_image.py
69
75
  tests/test_depth_adjusted_height_map.py
70
76
  tests/test_depth_estimate_height_map.py
77
+ tests/test_device_utils.py
71
78
  tests/test_end_to_end_optimizer.py
72
- tests/test_filament_and_output.py
79
+ tests/test_fasttsp_heightmap.py
73
80
  tests/test_filament_helper.py
74
81
  tests/test_flatforge.py
75
- tests/test_heightmap_init.py
76
82
  tests/test_image_helper.py
77
83
  tests/test_loss_functions.py
78
- tests/test_optimizer_helper.py
79
84
  tests/test_optimizer_module.py
80
- tests/test_optimizer_smoke.py
85
+ tests/test_optimizer_training.py
86
+ tests/test_other_helper.py
81
87
  tests/test_output_helper.py
82
- tests/test_output_helper_extended.py
83
- tests/test_perception_and_pruning.py
88
+ tests/test_perception_loss.py
89
+ tests/test_pipeline_smoke.py
84
90
  tests/test_property_based.py
91
+ tests/test_prune_repeatability.py
85
92
  tests/test_pruning_helper.py
86
- tests/test_remove_height_spikes.py
93
+ tests/test_remove_height_spikes.py
94
+ tests/test_webui_backend.py
95
+ tests/test_webui_regressions.py
@@ -20,3 +20,10 @@ python-multipart>=0.0.12
20
20
  aiofiles>=24.1.0
21
21
  websockets>=15.0
22
22
  pydantic-settings>=2.0.0
23
+ posthog>=3.7.0
24
+
25
+ [test]
26
+ pytest>=7.4
27
+ pytest-cov
28
+ hypothesis
29
+ httpx>=0.27
@@ -6,6 +6,8 @@ from typing import Optional, Tuple, Dict
6
6
  import os
7
7
  import torch
8
8
 
9
+ from autoforge.Helper.DeviceUtils import is_rocm
10
+
9
11
  # Cache the selection per device key so CPU/CUDA/MPS don't interfere
10
12
  _SELECTED_MAP: Dict[str, Tuple[Optional[torch.dtype], str]] = {}
11
13
 
@@ -20,11 +22,14 @@ def _device_key(device: torch.device) -> str:
20
22
  def _probe_backward(device: torch.device, dtype: torch.dtype) -> bool:
21
23
  """Try a tiny forward+backward under autocast for given dtype on the device.
22
24
 
23
- Returns True if successful, False if a RuntimeError occurs (e.g., oneDNN bf16/f16 backward not supported).
25
+ Returns True if successful, False if a RuntimeError occurs (e.g., oneDNN
26
+ bf16/f16 backward not supported, or an op with no reduced-precision Metal
27
+ kernel). The probe is the whole point: rather than maintaining a table of
28
+ which backend/driver/CPU-ISA combination supports what, we run the thing
29
+ and see. That is what makes CUDA, ROCm, Metal and CPU share one code path
30
+ here.
24
31
  """
25
32
  try:
26
- if device.type == "mps":
27
- return False # no stable autocast on MPS yet
28
33
  x = torch.randn(8, 8, device=device, dtype=torch.float32, requires_grad=False)
29
34
  w = torch.randn(8, 8, device=device, dtype=torch.float32, requires_grad=True)
30
35
  target = torch.randn(8, 8, device=device, dtype=torch.float32)
@@ -40,6 +45,33 @@ def _probe_backward(device: torch.device, dtype: torch.dtype) -> bool:
40
45
  return False
41
46
 
42
47
 
48
+ def _cuda_bf16_supported() -> bool:
49
+ try:
50
+ return bool(torch.cuda.is_bf16_supported())
51
+ except Exception:
52
+ return False
53
+
54
+
55
+ def _cuda_fp16_supported(device: torch.device) -> bool:
56
+ """fp16 capability check, tolerant of backends without CUDA arch numbers."""
57
+ if is_rocm():
58
+ return True
59
+ try:
60
+ major, minor = torch.cuda.get_device_capability(device)
61
+ except Exception:
62
+ return False
63
+ return (major * 10 + minor) >= 53 # CC 5.3 or newer
64
+
65
+
66
+ def _enable_tf32() -> None:
67
+ """Turn on TF32 matmuls where the backend has them (Ampere+; no-op on ROCm)."""
68
+ try:
69
+ torch.backends.cuda.matmul.allow_tf32 = True
70
+ torch.backends.cudnn.allow_tf32 = True
71
+ except Exception:
72
+ pass
73
+
74
+
43
75
  def _select_dtype_for_device(device: torch.device) -> Tuple[Optional[torch.dtype], str]:
44
76
  """Decide the best autocast dtype for the current device, probing backward support.
45
77
 
@@ -59,27 +91,30 @@ def _select_dtype_for_device(device: torch.device) -> Tuple[Optional[torch.dtype
59
91
  print("AUTOFORGE_AMP=bf16: disabling bf16 as probe failed")
60
92
  return None, "forced bf16 but probe failed"
61
93
  if mode in {"fp16", "float16", "half"}:
62
- if device.type == "cuda" and _probe_backward(device, torch.float16):
94
+ if device.type in {"cuda", "mps"} and _probe_backward(device, torch.float16):
63
95
  return torch.float16, "forced f16"
64
- print("AUTOFORGE_AMP=fp16: disabling fp16 as probe failed or non-CUDA")
65
- return None, "forced f16 but probe failed or non-CUDA"
96
+ print("AUTOFORGE_AMP=fp16: disabling fp16 as probe failed or unsupported device")
97
+ return None, "forced f16 but probe failed or unsupported device"
66
98
 
67
99
  # Automatic selection
68
100
  if device.type == "cuda":
101
+ # ROCm reports itself through the torch.cuda namespace, so the whole
102
+ # branch is shared; only the log label and the compute-capability
103
+ # lookup (which is CUDA-specific numbering) differ.
104
+ label = "ROCm" if is_rocm() else "CUDA"
69
105
  # Prefer bf16 if it actually works end-to-end
70
- if torch.cuda.is_bf16_supported() and _probe_backward(device, torch.bfloat16):
71
- print("Using bf16 autocast on CUDA device")
72
- return torch.bfloat16, "cuda bf16"
73
- # Next try fp16 if the GPU supports it
74
- major, minor = torch.cuda.get_device_capability(device)
75
- if (major * 10 + minor) >= 53 and _probe_backward(device, torch.float16):
76
- print("Using fp16 autocast on CUDA device")
77
- return torch.float16, "cuda f16"
106
+ if _cuda_bf16_supported() and _probe_backward(device, torch.bfloat16):
107
+ print(f"Using bf16 autocast on {label} device")
108
+ return torch.bfloat16, f"{label.lower()} bf16"
109
+ # Next try fp16 if the GPU supports it. Compute capability is a CUDA
110
+ # concept; on ROCm every supported GPU does fp16, so skip the gate.
111
+ if _cuda_fp16_supported(device) and _probe_backward(device, torch.float16):
112
+ print(f"Using fp16 autocast on {label} device")
113
+ return torch.float16, f"{label.lower()} f16"
78
114
  # Fall back to fp32; enable TF32 to speed up matmuls if available
79
- torch.backends.cuda.matmul.allow_tf32 = True
80
- torch.backends.cudnn.allow_tf32 = True
81
- print("Falling back to fp32 autocast on CUDA device with TF32 enabled")
82
- return None, "cuda fp32 (tf32 on)"
115
+ _enable_tf32()
116
+ print(f"Falling back to fp32 autocast on {label} device with TF32 enabled")
117
+ return None, f"{label.lower()} fp32 (tf32 on)"
83
118
 
84
119
  if device.type == "cpu":
85
120
  # Some CPUs (AVX512) can do bf16 forward/backward with oneDNN, but many cannot.
@@ -92,7 +127,11 @@ def _select_dtype_for_device(device: torch.device) -> Tuple[Optional[torch.dtype
92
127
  return None, "cpu fp32"
93
128
 
94
129
  if device.type == "mps":
95
- print("MPS device detected: autocast not reliably supported, using fp32")
130
+ # Metal does have an autocast implementation, but its reduced-precision
131
+ # coverage is uneven across ops and a missing kernel shows up as a
132
+ # wrong result rather than an error - so fp32 stays the default and
133
+ # AUTOFORGE_AMP=fp16/bf16 is the opt-in (handled above, probe-gated).
134
+ print("MPS device detected: defaulting to fp32 (set AUTOFORGE_AMP=fp16 to try half precision)")
96
135
  return None, "mps fp32"
97
136
 
98
137
  # Default: no autocast