autoforge 1.9.7__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 (125) hide show
  1. {autoforge-1.9.7/src/AutoForge.egg-info → autoforge-2.0.0}/PKG-INFO +35 -10
  2. autoforge-1.9.7/PKG-INFO → autoforge-2.0.0/README.md +27 -39
  3. {autoforge-1.9.7 → autoforge-2.0.0}/pyproject.toml +14 -2
  4. autoforge-1.9.7/README.md → autoforge-2.0.0/src/AutoForge.egg-info/PKG-INFO +64 -8
  5. {autoforge-1.9.7 → autoforge-2.0.0}/src/AutoForge.egg-info/SOURCES.txt +18 -9
  6. {autoforge-1.9.7 → autoforge-2.0.0}/src/AutoForge.egg-info/requires.txt +7 -0
  7. {autoforge-1.9.7 → 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.7 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/DepthEstimateHeightMap.py +21 -4
  10. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/FastTSPHeightMap.py +42 -24
  11. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/OptimizerHelper.py +94 -23
  12. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/OtherHelper.py +26 -7
  13. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Loss/LossFunctions.py +5 -3
  14. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Modules/Optimizer.py +464 -64
  15. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/auto_forge.py +41 -6
  16. autoforge-2.0.0/src/autoforge/webui/api/filaments.py +231 -0
  17. autoforge-2.0.0/src/autoforge/webui/api/images.py +49 -0
  18. autoforge-2.0.0/src/autoforge/webui/api/init.py +359 -0
  19. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/api/optimization.py +117 -22
  20. autoforge-2.0.0/src/autoforge/webui/api/outputs.py +212 -0
  21. autoforge-2.0.0/src/autoforge/webui/api/preview.py +157 -0
  22. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/api/project.py +6 -1
  23. autoforge-2.0.0/src/autoforge/webui/api/pruning.py +415 -0
  24. autoforge-2.0.0/src/autoforge/webui/api/sliders.py +107 -0
  25. autoforge-2.0.0/src/autoforge/webui/api/state.py +63 -0
  26. autoforge-2.0.0/src/autoforge/webui/api/system.py +189 -0
  27. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/api/ws.py +28 -6
  28. autoforge-2.0.0/src/autoforge/webui/config.py +74 -0
  29. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/helpers/colored_mesh.py +141 -21
  30. autoforge-2.0.0/src/autoforge/webui/helpers/gpu_memory.py +80 -0
  31. autoforge-2.0.0/src/autoforge/webui/helpers/mesh_persist.py +78 -0
  32. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/helpers/pipeline_runner.py +237 -54
  33. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/helpers/slider_render.py +152 -28
  34. autoforge-2.0.0/src/autoforge/webui/helpers/sliders.py +233 -0
  35. autoforge-2.0.0/src/autoforge/webui/helpers/telemetry.py +113 -0
  36. autoforge-2.0.0/src/autoforge/webui/models.py +219 -0
  37. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/server.py +14 -4
  38. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/services/filament_service.py +138 -41
  39. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/services/image_service.py +14 -2
  40. autoforge-2.0.0/src/autoforge/webui/services/optimization_service.py +440 -0
  41. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/services/project_service.py +39 -1
  42. autoforge-2.0.0/tests/test_amp_utils.py +94 -0
  43. autoforge-2.0.0/tests/test_auto_forge_cli.py +60 -0
  44. autoforge-2.0.0/tests/test_bug_regressions.py +413 -0
  45. autoforge-2.0.0/tests/test_cadamw.py +148 -0
  46. autoforge-2.0.0/tests/test_christofides_heightmap.py +183 -0
  47. autoforge-2.0.0/tests/test_composite_image.py +269 -0
  48. {autoforge-1.9.7 → autoforge-2.0.0}/tests/test_depth_adjusted_height_map.py +1 -1
  49. autoforge-2.0.0/tests/test_device_utils.py +287 -0
  50. autoforge-2.0.0/tests/test_fasttsp_heightmap.py +149 -0
  51. autoforge-2.0.0/tests/test_filament_helper.py +186 -0
  52. autoforge-2.0.0/tests/test_image_helper.py +118 -0
  53. autoforge-2.0.0/tests/test_loss_functions.py +166 -0
  54. autoforge-2.0.0/tests/test_optimizer_training.py +361 -0
  55. autoforge-2.0.0/tests/test_other_helper.py +81 -0
  56. autoforge-2.0.0/tests/test_output_helper.py +160 -0
  57. autoforge-2.0.0/tests/test_perception_loss.py +65 -0
  58. autoforge-2.0.0/tests/test_pipeline_smoke.py +209 -0
  59. autoforge-2.0.0/tests/test_prune_repeatability.py +375 -0
  60. autoforge-2.0.0/tests/test_pruning_helper.py +102 -0
  61. autoforge-2.0.0/tests/test_webui_backend.py +505 -0
  62. autoforge-2.0.0/tests/test_webui_regressions.py +351 -0
  63. autoforge-1.9.7/src/autoforge/webui/api/filaments.py +0 -128
  64. autoforge-1.9.7/src/autoforge/webui/api/images.py +0 -24
  65. autoforge-1.9.7/src/autoforge/webui/api/init.py +0 -37
  66. autoforge-1.9.7/src/autoforge/webui/api/outputs.py +0 -96
  67. autoforge-1.9.7/src/autoforge/webui/api/preview.py +0 -84
  68. autoforge-1.9.7/src/autoforge/webui/api/pruning.py +0 -175
  69. autoforge-1.9.7/src/autoforge/webui/api/sliders.py +0 -30
  70. autoforge-1.9.7/src/autoforge/webui/api/state.py +0 -38
  71. autoforge-1.9.7/src/autoforge/webui/api/system.py +0 -103
  72. autoforge-1.9.7/src/autoforge/webui/config.py +0 -27
  73. autoforge-1.9.7/src/autoforge/webui/helpers/sliders.py +0 -104
  74. autoforge-1.9.7/src/autoforge/webui/models.py +0 -118
  75. autoforge-1.9.7/src/autoforge/webui/services/optimization_service.py +0 -212
  76. autoforge-1.9.7/tests/test_amp_autocast_probe.py +0 -52
  77. autoforge-1.9.7/tests/test_cadamw.py +0 -50
  78. autoforge-1.9.7/tests/test_christofides_height_map.py +0 -63
  79. autoforge-1.9.7/tests/test_christofides_heightmap.py +0 -70
  80. autoforge-1.9.7/tests/test_filament_and_output.py +0 -105
  81. autoforge-1.9.7/tests/test_filament_helper.py +0 -75
  82. autoforge-1.9.7/tests/test_heightmap_init.py +0 -54
  83. autoforge-1.9.7/tests/test_image_helper.py +0 -42
  84. autoforge-1.9.7/tests/test_loss_functions.py +0 -33
  85. autoforge-1.9.7/tests/test_optimizer_helper.py +0 -65
  86. autoforge-1.9.7/tests/test_optimizer_smoke.py +0 -59
  87. autoforge-1.9.7/tests/test_output_helper.py +0 -105
  88. autoforge-1.9.7/tests/test_output_helper_extended.py +0 -88
  89. autoforge-1.9.7/tests/test_perception_and_pruning.py +0 -103
  90. autoforge-1.9.7/tests/test_pruning_helper.py +0 -89
  91. {autoforge-1.9.7 → autoforge-2.0.0}/setup.cfg +0 -0
  92. {autoforge-1.9.7 → autoforge-2.0.0}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  93. {autoforge-1.9.7 → autoforge-2.0.0}/src/AutoForge.egg-info/entry_points.txt +0 -0
  94. {autoforge-1.9.7 → autoforge-2.0.0}/src/AutoForge.egg-info/top_level.txt +0 -0
  95. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/CAdamW.py +0 -0
  96. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/FilamentHelper.py +0 -0
  97. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/ChristofidesHeightMap.py +0 -0
  98. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/Heightmaps/__init__.py +0 -0
  99. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/ImageHelper.py +0 -0
  100. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/OutputHelper.py +0 -0
  101. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/PruningHelper.py +0 -0
  102. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/__init__.py +0 -0
  103. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Helper/filamentcolors_library.py +0 -0
  104. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Init_comparer/__init__.py +0 -0
  105. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Init_comparer/auto_forge_init_check.py +0 -0
  106. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Init_comparer/plot_results.py +0 -0
  107. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  108. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Loss/__init__.py +0 -0
  109. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Misc/__init__.py +0 -0
  110. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Misc/color_fit_all.py +0 -0
  111. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/Modules/__init__.py +0 -0
  112. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/__init__.py +0 -0
  113. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/__main__.py +0 -0
  114. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/priority_mask_editor.py +0 -0
  115. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/__init__.py +0 -0
  116. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/api/__init__.py +0 -0
  117. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/api/settings.py +0 -0
  118. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/helpers/__init__.py +0 -0
  119. {autoforge-1.9.7 → autoforge-2.0.0}/src/autoforge/webui/services/__init__.py +0 -0
  120. {autoforge-1.9.7 → autoforge-2.0.0}/tests/test_depth_estimate_height_map.py +0 -0
  121. {autoforge-1.9.7 → autoforge-2.0.0}/tests/test_end_to_end_optimizer.py +0 -0
  122. {autoforge-1.9.7 → autoforge-2.0.0}/tests/test_flatforge.py +0 -0
  123. {autoforge-1.9.7 → autoforge-2.0.0}/tests/test_optimizer_module.py +0 -0
  124. {autoforge-1.9.7 → autoforge-2.0.0}/tests/test_property_based.py +0 -0
  125. {autoforge-1.9.7 → 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.7
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.
@@ -74,19 +79,27 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
74
79
 
75
80
  ## Web UI: One-Click Install & Run
76
81
 
77
- The easiest way to use AutoForge is the web UI — a local app (like ComfyUI) with drag-and-drop image upload, a filament library, live sliders for adjusting colors after optimization, and pruning. No command-line arguments needed.
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.
78
84
 
79
85
  1. **Clone this repository** (or download and extract the ZIP from the green "Code" button on GitHub).
80
86
  2. **Install**, from the project folder:
81
87
  - Linux/macOS: `./install.sh`
82
88
  - Windows: double-click `install.bat` (or run it from a terminal)
83
89
 
84
- 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.
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.
85
91
  3. **Run**:
86
92
  - Linux/macOS: `./run_webui.sh`
87
93
  - Windows: double-click `run_webui.bat`
88
94
 
89
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.
100
+
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
+
90
103
  4. **Update to the latest release** whenever you want, from the project folder:
91
104
  - Linux/macOS: `./update.sh`
92
105
  - Windows: double-click `update.bat`
@@ -94,7 +107,7 @@ The easiest way to use AutoForge is the web UI — a local app (like ComfyUI) wi
94
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`).
95
108
 
96
109
  If you have problems running the code on your GPU, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
97
- 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.
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.)
98
111
 
99
112
  ## Manual Installation (CLI only)
100
113
 
@@ -170,6 +183,8 @@ This will generate separate STL files for each color, allowing you to print face
170
183
  - `--nozzle_diameter` Diameter of the printer nozzle in millimeters (default: 0.4).
171
184
  **Note:** Details smaller than half this value will be ignored.
172
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.
173
188
 
174
189
  - `--flatforge` Enable FlatForge mode to generate separate STL files for each color (default: False).
175
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).
@@ -187,7 +202,8 @@ This will generate separate STL files for each color, allowing you to print face
187
202
  - `--pruning_max_swaps` Max number of swaps allowed after pruning (default: 100).
188
203
  - `--pruning_max_layer` Max number of layers allowed after pruning (default: 75).
189
204
  - `--random_seed` Random seed for reproducibility (default: 0 (disabled)).
190
- - `--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`.
191
207
  - `--no-spike-removal` Disable spike removal for the final STL (not recommended).
192
208
 
193
209
  - `--tensorboard` Flag to enable TensorBoard logging.
@@ -216,6 +232,10 @@ After running, the following files will be created in your specified output fold
216
232
 
217
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.
218
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
+
219
239
  ## Development
220
240
 
221
241
  To have a "nightly" version of the repository or have live updating changes during development please do the following:
@@ -245,6 +265,11 @@ pip install torch torchvision torchaudio --index-url https://download.pytorch.or
245
265
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
246
266
  The software is provided as-is and comes with no warranty or guarantee of support.
247
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.
248
273
 
249
274
  ## Acknowledgements
250
275
 
@@ -1,44 +1,12 @@
1
- Metadata-Version: 2.4
2
- Name: AutoForge
3
- Version: 1.9.7
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.
@@ -74,19 +42,27 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
74
42
 
75
43
  ## Web UI: One-Click Install & Run
76
44
 
77
- The easiest way to use AutoForge is the web UI — a local app (like ComfyUI) with drag-and-drop image upload, a filament library, live sliders for adjusting colors after optimization, and pruning. No command-line arguments needed.
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.
78
47
 
79
48
  1. **Clone this repository** (or download and extract the ZIP from the green "Code" button on GitHub).
80
49
  2. **Install**, from the project folder:
81
50
  - Linux/macOS: `./install.sh`
82
51
  - Windows: double-click `install.bat` (or run it from a terminal)
83
52
 
84
- 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.
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.
85
54
  3. **Run**:
86
55
  - Linux/macOS: `./run_webui.sh`
87
56
  - Windows: double-click `run_webui.bat`
88
57
 
89
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.
63
+
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
+
90
66
  4. **Update to the latest release** whenever you want, from the project folder:
91
67
  - Linux/macOS: `./update.sh`
92
68
  - Windows: double-click `update.bat`
@@ -94,7 +70,7 @@ The easiest way to use AutoForge is the web UI — a local app (like ComfyUI) wi
94
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`).
95
71
 
96
72
  If you have problems running the code on your GPU, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
97
- 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.
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.)
98
74
 
99
75
  ## Manual Installation (CLI only)
100
76
 
@@ -170,6 +146,8 @@ This will generate separate STL files for each color, allowing you to print face
170
146
  - `--nozzle_diameter` Diameter of the printer nozzle in millimeters (default: 0.4).
171
147
  **Note:** Details smaller than half this value will be ignored.
172
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.
173
151
 
174
152
  - `--flatforge` Enable FlatForge mode to generate separate STL files for each color (default: False).
175
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).
@@ -187,7 +165,8 @@ This will generate separate STL files for each color, allowing you to print face
187
165
  - `--pruning_max_swaps` Max number of swaps allowed after pruning (default: 100).
188
166
  - `--pruning_max_layer` Max number of layers allowed after pruning (default: 75).
189
167
  - `--random_seed` Random seed for reproducibility (default: 0 (disabled)).
190
- - `--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`.
191
170
  - `--no-spike-removal` Disable spike removal for the final STL (not recommended).
192
171
 
193
172
  - `--tensorboard` Flag to enable TensorBoard logging.
@@ -216,6 +195,10 @@ After running, the following files will be created in your specified output fold
216
195
 
217
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.
218
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
+
219
202
  ## Development
220
203
 
221
204
  To have a "nightly" version of the repository or have live updating changes during development please do the following:
@@ -245,6 +228,11 @@ pip install torch torchvision torchaudio --index-url https://download.pytorch.or
245
228
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
246
229
  The software is provided as-is and comes with no warranty or guarantee of support.
247
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.
248
236
 
249
237
  ## Acknowledgements
250
238
 
@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="1.9.7"
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.
@@ -43,19 +79,27 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
43
79
 
44
80
  ## Web UI: One-Click Install & Run
45
81
 
46
- The easiest way to use AutoForge is the web UI — a local app (like ComfyUI) with drag-and-drop image upload, a filament library, live sliders for adjusting colors after optimization, and pruning. No command-line arguments needed.
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.
47
84
 
48
85
  1. **Clone this repository** (or download and extract the ZIP from the green "Code" button on GitHub).
49
86
  2. **Install**, from the project folder:
50
87
  - Linux/macOS: `./install.sh`
51
88
  - Windows: double-click `install.bat` (or run it from a terminal)
52
89
 
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.
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.
54
91
  3. **Run**:
55
92
  - Linux/macOS: `./run_webui.sh`
56
93
  - Windows: double-click `run_webui.bat`
57
94
 
58
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.
100
+
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
+
59
103
  4. **Update to the latest release** whenever you want, from the project folder:
60
104
  - Linux/macOS: `./update.sh`
61
105
  - Windows: double-click `update.bat`
@@ -63,7 +107,7 @@ The easiest way to use AutoForge is the web UI — a local app (like ComfyUI) wi
63
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`).
64
108
 
65
109
  If you have problems running the code on your GPU, please refer to the [Pytorch Homepage](https://pytorch.org/) for help. \
66
- 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.
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.)
67
111
 
68
112
  ## Manual Installation (CLI only)
69
113
 
@@ -139,6 +183,8 @@ This will generate separate STL files for each color, allowing you to print face
139
183
  - `--nozzle_diameter` Diameter of the printer nozzle in millimeters (default: 0.4).
140
184
  **Note:** Details smaller than half this value will be ignored.
141
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.
142
188
 
143
189
  - `--flatforge` Enable FlatForge mode to generate separate STL files for each color (default: False).
144
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).
@@ -156,7 +202,8 @@ This will generate separate STL files for each color, allowing you to print face
156
202
  - `--pruning_max_swaps` Max number of swaps allowed after pruning (default: 100).
157
203
  - `--pruning_max_layer` Max number of layers allowed after pruning (default: 75).
158
204
  - `--random_seed` Random seed for reproducibility (default: 0 (disabled)).
159
- - `--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`.
160
207
  - `--no-spike-removal` Disable spike removal for the final STL (not recommended).
161
208
 
162
209
  - `--tensorboard` Flag to enable TensorBoard logging.
@@ -185,6 +232,10 @@ After running, the following files will be created in your specified output fold
185
232
 
186
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.
187
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
+
188
239
  ## Development
189
240
 
190
241
  To have a "nightly" version of the repository or have live updating changes during development please do the following:
@@ -214,6 +265,11 @@ pip install torch torchvision torchaudio --index-url https://download.pytorch.or
214
265
  AutoForge © 2025 by Hendric Voss is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).
215
266
  The software is provided as-is and comes with no warranty or guarantee of support.
216
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.
217
273
 
218
274
  ## Acknowledgements
219
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