autoforge 2.0.0__tar.gz → 2.1.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 (111) hide show
  1. {autoforge-2.0.0 → autoforge-2.1.0}/PKG-INFO +7 -2
  2. {autoforge-2.0.0 → autoforge-2.1.0}/README.md +5 -0
  3. {autoforge-2.0.0 → autoforge-2.1.0}/pyproject.toml +5 -2
  4. {autoforge-2.0.0 → autoforge-2.1.0}/src/AutoForge.egg-info/PKG-INFO +7 -2
  5. {autoforge-2.0.0 → autoforge-2.1.0}/src/AutoForge.egg-info/SOURCES.txt +12 -0
  6. autoforge-2.1.0/src/autoforge/Helper/ConstraintHelper.py +279 -0
  7. autoforge-2.1.0/src/autoforge/Helper/FusedComposite.py +900 -0
  8. autoforge-2.1.0/src/autoforge/Helper/HeightAssign.py +148 -0
  9. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/Heightmaps/FastTSPHeightMap.py +19 -3
  10. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/ImageHelper.py +29 -0
  11. autoforge-2.1.0/src/autoforge/Helper/OptimizerHelper.py +1432 -0
  12. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/OtherHelper.py +4 -2
  13. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/OutputHelper.py +135 -145
  14. autoforge-2.1.0/src/autoforge/Helper/PixelHeightRefine.py +1053 -0
  15. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/PruningHelper.py +299 -92
  16. autoforge-2.1.0/src/autoforge/Helper/filamentcolors_library.py +345 -0
  17. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Modules/Optimizer.py +930 -52
  18. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/auto_forge.py +393 -25
  19. autoforge-2.1.0/src/autoforge/data/filamentcolors_catalog.json +10216 -0
  20. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/filaments.py +49 -0
  21. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/init.py +13 -3
  22. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/optimization.py +73 -5
  23. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/preview.py +8 -2
  24. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/pruning.py +92 -18
  25. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/ws.py +5 -0
  26. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/config.py +6 -0
  27. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/pipeline_runner.py +87 -36
  28. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/slider_render.py +29 -19
  29. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/sliders.py +15 -7
  30. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/models.py +5 -0
  31. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/server.py +7 -0
  32. autoforge-2.1.0/src/autoforge/webui/services/catalog_service.py +150 -0
  33. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/services/filament_service.py +60 -0
  34. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/services/optimization_service.py +15 -0
  35. autoforge-2.1.0/tests/test_bugfixes_2026_09_26.py +414 -0
  36. autoforge-2.1.0/tests/test_constrained_opt.py +596 -0
  37. autoforge-2.1.0/tests/test_filamentcolors_library.py +345 -0
  38. autoforge-2.1.0/tests/test_graph_cache_eviction.py +117 -0
  39. autoforge-2.1.0/tests/test_graph_capture_threads.py +53 -0
  40. autoforge-2.1.0/tests/test_gumbel_nan.py +25 -0
  41. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_optimizer_training.py +6 -2
  42. autoforge-2.0.0/src/autoforge/Helper/OptimizerHelper.py +0 -822
  43. autoforge-2.0.0/src/autoforge/Helper/filamentcolors_library.py +0 -164
  44. {autoforge-2.0.0 → autoforge-2.1.0}/setup.cfg +0 -0
  45. {autoforge-2.0.0 → autoforge-2.1.0}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  46. {autoforge-2.0.0 → autoforge-2.1.0}/src/AutoForge.egg-info/entry_points.txt +0 -0
  47. {autoforge-2.0.0 → autoforge-2.1.0}/src/AutoForge.egg-info/requires.txt +0 -0
  48. {autoforge-2.0.0 → autoforge-2.1.0}/src/AutoForge.egg-info/top_level.txt +0 -0
  49. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/AmpUtils.py +0 -0
  50. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/CAdamW.py +0 -0
  51. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/DeviceUtils.py +0 -0
  52. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/FilamentHelper.py +0 -0
  53. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/Heightmaps/ChristofidesHeightMap.py +0 -0
  54. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/Heightmaps/DepthEstimateHeightMap.py +0 -0
  55. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/Heightmaps/__init__.py +0 -0
  56. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Helper/__init__.py +0 -0
  57. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Init_comparer/__init__.py +0 -0
  58. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Init_comparer/auto_forge_init_check.py +0 -0
  59. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Init_comparer/plot_results.py +0 -0
  60. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Loss/LossFunctions.py +0 -0
  61. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  62. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Loss/__init__.py +0 -0
  63. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Misc/__init__.py +0 -0
  64. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Misc/color_fit_all.py +0 -0
  65. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/Modules/__init__.py +0 -0
  66. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/__init__.py +0 -0
  67. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/__main__.py +0 -0
  68. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/priority_mask_editor.py +0 -0
  69. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/__init__.py +0 -0
  70. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/__init__.py +0 -0
  71. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/images.py +0 -0
  72. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/outputs.py +0 -0
  73. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/project.py +0 -0
  74. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/settings.py +0 -0
  75. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/sliders.py +0 -0
  76. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/state.py +0 -0
  77. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/api/system.py +0 -0
  78. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/__init__.py +0 -0
  79. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/colored_mesh.py +0 -0
  80. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/gpu_memory.py +0 -0
  81. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/mesh_persist.py +0 -0
  82. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/helpers/telemetry.py +0 -0
  83. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/services/__init__.py +0 -0
  84. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/services/image_service.py +0 -0
  85. {autoforge-2.0.0 → autoforge-2.1.0}/src/autoforge/webui/services/project_service.py +0 -0
  86. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_amp_utils.py +0 -0
  87. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_auto_forge_cli.py +0 -0
  88. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_bug_regressions.py +0 -0
  89. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_cadamw.py +0 -0
  90. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_christofides_heightmap.py +0 -0
  91. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_composite_image.py +0 -0
  92. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_depth_adjusted_height_map.py +0 -0
  93. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_depth_estimate_height_map.py +0 -0
  94. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_device_utils.py +0 -0
  95. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_end_to_end_optimizer.py +0 -0
  96. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_fasttsp_heightmap.py +0 -0
  97. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_filament_helper.py +0 -0
  98. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_flatforge.py +0 -0
  99. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_image_helper.py +0 -0
  100. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_loss_functions.py +0 -0
  101. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_optimizer_module.py +0 -0
  102. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_other_helper.py +0 -0
  103. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_output_helper.py +0 -0
  104. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_perception_loss.py +0 -0
  105. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_pipeline_smoke.py +0 -0
  106. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_property_based.py +0 -0
  107. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_prune_repeatability.py +0 -0
  108. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_pruning_helper.py +0 -0
  109. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_remove_height_spikes.py +0 -0
  110. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_webui_backend.py +0 -0
  111. {autoforge-2.0.0 → autoforge-2.1.0}/tests/test_webui_regressions.py +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: AutoForge
3
- Version: 2.0.0
4
- Summary: AutoForge - 3D color images made easy
3
+ Version: 2.1.0
4
+ Summary: AutoForge - 3D color prints 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
@@ -98,6 +98,8 @@ The easiest way to use AutoForge is the web UI, a local app with drag-and-drop i
98
98
 
99
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
100
 
101
+ To find a filament you own without measuring it yourself, click **Catalog** in the filament library. It searches every filament on [filamentcolors.xyz](https://filamentcolors.xyz) that has a measured TD (search by brand, color name, type or hex code; filter by type, brand or color family; or pick a color to see the closest matches first), and **Add** puts it straight into your library.
102
+
101
103
  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
104
 
103
105
  4. **Update to the latest release** whenever you want, from the project folder:
@@ -248,6 +250,8 @@ conda activate forge
248
250
  pip install -e .
249
251
  ```
250
252
 
253
+ To refresh the bundled filamentcolors.xyz catalog (`src/autoforge/data/filamentcolors_catalog.json`) before a release, run `./generate_filamentcolors_library.sh`. It downloads the whole catalog once, pausing 2 seconds between pages (`--delay` sets the pause; requests are never less than 1 second apart), and keeps only the filaments that have a TD.
254
+
251
255
  If the installed pytorch version has no cuda support execute the following:
252
256
 
253
257
  ```bash
@@ -276,6 +280,7 @@ I would love to see what you have done with the software, so it would be great i
276
280
  First and foremost:
277
281
  - [Hueforge](https://shop.thehueforge.com/) for providing the inspiration for this project.
278
282
  Without it, this project would not have been possible.
283
+ - [filamentcolors.xyz](https://filamentcolors.xyz) and its community for the measured filament swatches and TD values in the web UI's filament catalog.
279
284
 
280
285
  AutoForge makes use of several open source libraries:
281
286
 
@@ -61,6 +61,8 @@ The easiest way to use AutoForge is the web UI, a local app with drag-and-drop i
61
61
 
62
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
63
 
64
+ To find a filament you own without measuring it yourself, click **Catalog** in the filament library. It searches every filament on [filamentcolors.xyz](https://filamentcolors.xyz) that has a measured TD (search by brand, color name, type or hex code; filter by type, brand or color family; or pick a color to see the closest matches first), and **Add** puts it straight into your library.
65
+
64
66
  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
67
 
66
68
  4. **Update to the latest release** whenever you want, from the project folder:
@@ -211,6 +213,8 @@ conda activate forge
211
213
  pip install -e .
212
214
  ```
213
215
 
216
+ To refresh the bundled filamentcolors.xyz catalog (`src/autoforge/data/filamentcolors_catalog.json`) before a release, run `./generate_filamentcolors_library.sh`. It downloads the whole catalog once, pausing 2 seconds between pages (`--delay` sets the pause; requests are never less than 1 second apart), and keeps only the filaments that have a TD.
217
+
214
218
  If the installed pytorch version has no cuda support execute the following:
215
219
 
216
220
  ```bash
@@ -239,6 +243,7 @@ I would love to see what you have done with the software, so it would be great i
239
243
  First and foremost:
240
244
  - [Hueforge](https://shop.thehueforge.com/) for providing the inspiration for this project.
241
245
  Without it, this project would not have been possible.
246
+ - [filamentcolors.xyz](https://filamentcolors.xyz) and its community for the measured filament swatches and TD values in the web UI's filament catalog.
242
247
 
243
248
  AutoForge makes use of several open source libraries:
244
249
 
@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="2.0.0"
8
- description = "AutoForge - 3D color images made easy"
7
+ version="2.1.0"
8
+ description = "AutoForge - 3D color prints 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"
@@ -50,6 +50,9 @@ test = [
50
50
  [tool.setuptools]
51
51
  package-dir = {"" = "src"}
52
52
 
53
+ [tool.setuptools.package-data]
54
+ autoforge = ["data/*.json"]
55
+
53
56
  [project.scripts]
54
57
  autoforge = "autoforge.auto_forge:main"
55
58
 
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: AutoForge
3
- Version: 2.0.0
4
- Summary: AutoForge - 3D color images made easy
3
+ Version: 2.1.0
4
+ Summary: AutoForge - 3D color prints 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
@@ -98,6 +98,8 @@ The easiest way to use AutoForge is the web UI, a local app with drag-and-drop i
98
98
 
99
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
100
 
101
+ To find a filament you own without measuring it yourself, click **Catalog** in the filament library. It searches every filament on [filamentcolors.xyz](https://filamentcolors.xyz) that has a measured TD (search by brand, color name, type or hex code; filter by type, brand or color family; or pick a color to see the closest matches first), and **Add** puts it straight into your library.
102
+
101
103
  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
104
 
103
105
  4. **Update to the latest release** whenever you want, from the project folder:
@@ -248,6 +250,8 @@ conda activate forge
248
250
  pip install -e .
249
251
  ```
250
252
 
253
+ To refresh the bundled filamentcolors.xyz catalog (`src/autoforge/data/filamentcolors_catalog.json`) before a release, run `./generate_filamentcolors_library.sh`. It downloads the whole catalog once, pausing 2 seconds between pages (`--delay` sets the pause; requests are never less than 1 second apart), and keeps only the filaments that have a TD.
254
+
251
255
  If the installed pytorch version has no cuda support execute the following:
252
256
 
253
257
  ```bash
@@ -276,6 +280,7 @@ I would love to see what you have done with the software, so it would be great i
276
280
  First and foremost:
277
281
  - [Hueforge](https://shop.thehueforge.com/) for providing the inspiration for this project.
278
282
  Without it, this project would not have been possible.
283
+ - [filamentcolors.xyz](https://filamentcolors.xyz) and its community for the measured filament swatches and TD values in the web UI's filament catalog.
279
284
 
280
285
  AutoForge makes use of several open source libraries:
281
286
 
@@ -12,12 +12,16 @@ 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/ConstraintHelper.py
15
16
  src/autoforge/Helper/DeviceUtils.py
16
17
  src/autoforge/Helper/FilamentHelper.py
18
+ src/autoforge/Helper/FusedComposite.py
19
+ src/autoforge/Helper/HeightAssign.py
17
20
  src/autoforge/Helper/ImageHelper.py
18
21
  src/autoforge/Helper/OptimizerHelper.py
19
22
  src/autoforge/Helper/OtherHelper.py
20
23
  src/autoforge/Helper/OutputHelper.py
24
+ src/autoforge/Helper/PixelHeightRefine.py
21
25
  src/autoforge/Helper/PruningHelper.py
22
26
  src/autoforge/Helper/__init__.py
23
27
  src/autoforge/Helper/filamentcolors_library.py
@@ -35,6 +39,7 @@ src/autoforge/Misc/__init__.py
35
39
  src/autoforge/Misc/color_fit_all.py
36
40
  src/autoforge/Modules/Optimizer.py
37
41
  src/autoforge/Modules/__init__.py
42
+ src/autoforge/data/filamentcolors_catalog.json
38
43
  src/autoforge/webui/__init__.py
39
44
  src/autoforge/webui/config.py
40
45
  src/autoforge/webui/models.py
@@ -62,6 +67,7 @@ src/autoforge/webui/helpers/slider_render.py
62
67
  src/autoforge/webui/helpers/sliders.py
63
68
  src/autoforge/webui/helpers/telemetry.py
64
69
  src/autoforge/webui/services/__init__.py
70
+ src/autoforge/webui/services/catalog_service.py
65
71
  src/autoforge/webui/services/filament_service.py
66
72
  src/autoforge/webui/services/image_service.py
67
73
  src/autoforge/webui/services/optimization_service.py
@@ -69,16 +75,22 @@ src/autoforge/webui/services/project_service.py
69
75
  tests/test_amp_utils.py
70
76
  tests/test_auto_forge_cli.py
71
77
  tests/test_bug_regressions.py
78
+ tests/test_bugfixes_2026_09_26.py
72
79
  tests/test_cadamw.py
73
80
  tests/test_christofides_heightmap.py
74
81
  tests/test_composite_image.py
82
+ tests/test_constrained_opt.py
75
83
  tests/test_depth_adjusted_height_map.py
76
84
  tests/test_depth_estimate_height_map.py
77
85
  tests/test_device_utils.py
78
86
  tests/test_end_to_end_optimizer.py
79
87
  tests/test_fasttsp_heightmap.py
80
88
  tests/test_filament_helper.py
89
+ tests/test_filamentcolors_library.py
81
90
  tests/test_flatforge.py
91
+ tests/test_graph_cache_eviction.py
92
+ tests/test_graph_capture_threads.py
93
+ tests/test_gumbel_nan.py
82
94
  tests/test_image_helper.py
83
95
  tests/test_loss_functions.py
84
96
  tests/test_optimizer_module.py
@@ -0,0 +1,279 @@
1
+ """Colour / swap limits held by the optimizer itself (--constrained_opt).
2
+
3
+ ``project_assignment`` maps per-layer material scores [L, M] to the best
4
+ discrete layer stack that uses at most ``max_colors`` distinct materials
5
+ besides the base filament (``free``: layers may reuse it at no cost) and
6
+ at most ``max_swaps`` material changes between consecutive layers, where
7
+ "best" means the highest summed per-layer score. With log-probabilities as
8
+ scores that is the most likely feasible stack; with Gumbel-perturbed
9
+ log-probabilities it is a random feasible sample (perturb-and-MAP), which is
10
+ what the discrete snapshots during training use.
11
+ """
12
+ from itertools import combinations
13
+ from typing import Optional
14
+
15
+ import numpy as np
16
+ import torch
17
+
18
+ # Enumerate every palette up to this many subsets, greedy search beyond.
19
+ MAX_SUBSETS = 5000
20
+
21
+
22
+ def _mat(base: Optional[int]) -> Optional[int]:
23
+ """The base's filament index (None when there is no base filament)."""
24
+ return base if base is not None and base >= 0 else None
25
+
26
+
27
+ def _viterbi_swaps(score: torch.Tensor, max_swaps: int, base: Optional[int] = None) -> torch.Tensor:
28
+ """argmax over stacks with <= max_swaps changes of sum_l score[b, l, dg[l]].
29
+
30
+ With ``base`` (the base's filament, -1 for a base color that is no
31
+ filament) the change from the base to the first layer counts too: the
32
+ base is the stack's band 0.
33
+
34
+ score: [B, L, M] (use a large negative score for forbidden materials).
35
+ Returns [B, L] int64 on score's device.
36
+
37
+ Runs in NumPy on the CPU: the tables are tiny ([B, max_swaps+1, M] per
38
+ layer), and as a loop of small GPU kernels it took ~18ms per call with
39
+ the GPU idle in between - the bulk of the constrained snapshot cost.
40
+ """
41
+ dev = score.device
42
+ # A NaN score makes every comparison false, which backtracked into a
43
+ # swap count of -1.
44
+ sc = np.nan_to_num(
45
+ score.detach().to("cpu", torch.float64).numpy(), nan=-1e9, posinf=1e9, neginf=-1e9
46
+ )
47
+ B, L, M = sc.shape
48
+ S = max_swaps + 1
49
+ neg = -1e300
50
+ dp = np.full((B, S, M), neg)
51
+ if base is None:
52
+ dp[:, 0] = sc[:, 0]
53
+ else:
54
+ # Staying in the base filament is free, anything else is a swap.
55
+ if 0 <= base < M:
56
+ dp[:, 0, base] = sc[:, 0, base]
57
+ if S > 1:
58
+ other = np.arange(M) != base
59
+ dp[:, 1, other] = sc[:, 0, other]
60
+ back_stay = np.empty((max(L - 1, 0), B, S, M), dtype=bool)
61
+ back_from = np.empty((max(L - 1, 0), B, S, M), dtype=np.int64)
62
+ ar = np.arange(M)
63
+ for l in range(1, L):
64
+ # Best previous material other than m, one swap earlier.
65
+ i1 = dp.argmax(axis=2) # [B,S]
66
+ v1 = np.take_along_axis(dp, i1[..., None], axis=2)[..., 0]
67
+ masked = dp.copy()
68
+ np.put_along_axis(masked, i1[..., None], neg, axis=2)
69
+ i2 = masked.argmax(axis=2)
70
+ v2 = np.take_along_axis(masked, i2[..., None], axis=2)[..., 0]
71
+ is_best = i1[..., None] == ar
72
+ other_v = np.where(is_best, v2[..., None], v1[..., None]) # [B,S,M]
73
+ other_i = np.where(is_best, i2[..., None], i1[..., None])
74
+ switch_v = np.full_like(dp, neg)
75
+ switch_v[:, 1:] = other_v[:, :-1]
76
+ switch_i = np.zeros_like(other_i)
77
+ switch_i[:, 1:] = other_i[:, :-1]
78
+ stay = dp >= switch_v
79
+ stay[:, 0] = True # no swap left to come from
80
+ dp = np.where(stay, dp, switch_v) + sc[:, l][:, None, :]
81
+ back_stay[l - 1] = stay
82
+ back_from[l - 1] = switch_i
83
+ flat = dp.reshape(B, -1).argmax(axis=1)
84
+ s_idx, m_idx = flat // M, flat % M
85
+ out = np.empty((B, L), dtype=np.int64)
86
+ out[:, L - 1] = m_idx
87
+ bi = np.arange(B)
88
+ for l in range(L - 1, 0, -1):
89
+ st = back_stay[l - 1][bi, s_idx, m_idx]
90
+ prev = np.where(st, m_idx, back_from[l - 1][bi, s_idx, m_idx])
91
+ s_idx = np.where(st, s_idx, s_idx - 1)
92
+ m_idx = prev
93
+ out[:, l - 1] = m_idx
94
+ return torch.from_numpy(out).to(dev)
95
+
96
+
97
+ def _palette_candidates(M: int, max_colors: int, free: Optional[int], device) -> Optional[torch.Tensor]:
98
+ """All palettes as [N, M] bool: ``max_colors`` materials besides
99
+ ``free`` (the base filament, always in the palette at no cost), or None
100
+ when there are too many to enumerate."""
101
+ others = [m for m in range(M) if m != free]
102
+ k = min(max_colors, len(others))
103
+ n_sub = 1
104
+ for i in range(k):
105
+ n_sub = n_sub * (len(others) - i) // (i + 1)
106
+ if n_sub > MAX_SUBSETS:
107
+ return None
108
+ subs = torch.tensor(list(combinations(others, k)), device=device, dtype=torch.long).view(-1, k)
109
+ masks = torch.zeros((subs.shape[0], M), dtype=torch.bool, device=device)
110
+ masks.scatter_(1, subs, True)
111
+ if free is not None:
112
+ masks[:, free] = True
113
+ return masks
114
+
115
+
116
+ def top_palettes(score: torch.Tensor, max_colors: int, k: int, free: Optional[int] = None) -> torch.Tensor:
117
+ """The k best palettes for scores [L, M] by sum_l max_{m in palette}
118
+ score[l, m], as [k', M] bool (k' <= k; exhaustive search only)."""
119
+ L, M = score.shape
120
+ free = _mat(free)
121
+ masks = _palette_candidates(M, max_colors, free, score.device)
122
+ if masks is None:
123
+ return _palette_masks(score.unsqueeze(0), max_colors, free)
124
+ neg = torch.tensor(torch.finfo(score.dtype).min / 4, device=score.device)
125
+ vals = torch.where(masks.unsqueeze(1), score.unsqueeze(0), neg).amax(-1).sum(-1)
126
+ return masks[vals.topk(min(k, vals.shape[0])).indices]
127
+
128
+
129
+ def _palette_masks(score: torch.Tensor, max_colors: int, free: Optional[int] = None) -> torch.Tensor:
130
+ """Best palette per batch item, as [B, M] bool: at most ``max_colors``
131
+ materials besides ``free`` (the base filament, which is always in it),
132
+ judged by sum_l max_{m in palette} score[b, l, m]."""
133
+ B, L, M = score.shape
134
+ dev = score.device
135
+ neg = torch.finfo(score.dtype).min / 4
136
+ masks = _palette_candidates(M, max_colors, free, dev)
137
+ if masks is not None:
138
+ # [B, N, L]: best in-palette score per layer
139
+ vals = torch.where(
140
+ masks.view(1, -1, 1, M), score.unsqueeze(1), torch.tensor(neg, device=dev)
141
+ ).amax(dim=-1)
142
+ best = vals.sum(-1).argmax(dim=1) # [B]
143
+ return masks[best]
144
+ # Greedy: drop the material whose removal costs least, until small enough.
145
+ mask = torch.ones((B, M), dtype=torch.bool, device=dev)
146
+ n_other = M - (1 if free is not None else 0)
147
+ for _ in range(max(0, n_other - max_colors)):
148
+ costs = []
149
+ for m in range(M):
150
+ trial = mask.clone()
151
+ trial[:, m] = False
152
+ v = torch.where(trial.unsqueeze(1), score, torch.tensor(neg, device=dev))
153
+ c = v.amax(-1).sum(-1)
154
+ keep = mask[:, m] & (m != free)
155
+ c = torch.where(keep, c, torch.full_like(c, neg))
156
+ costs.append(c)
157
+ drop = torch.stack(costs, 1).argmax(1)
158
+ mask[torch.arange(B, device=dev), drop] = False
159
+ return mask
160
+
161
+
162
+ def project_assignment(
163
+ score: torch.Tensor,
164
+ max_colors: Optional[int],
165
+ max_swaps: Optional[int],
166
+ free: Optional[int] = None,
167
+ ) -> torch.Tensor:
168
+ """Best feasible stack for scores [L, M] or [B, L, M] -> [L] / [B, L].
169
+
170
+ ``free`` is the base: its filament index (a layer in it adds no colour
171
+ to the print, and a first layer in it is no swap), -1 for a base colour
172
+ that is no filament (a colour of its own, and the first layer is always
173
+ a swap), or None to leave the base out entirely. ``max_colors`` limits
174
+ the materials besides the base filament."""
175
+ squeeze = score.dim() == 2
176
+ if squeeze:
177
+ score = score.unsqueeze(0)
178
+ score = score.float()
179
+ M = score.shape[-1]
180
+ mat = _mat(free)
181
+ if max_colors is not None and max_colors < M - (1 if mat is not None else 0):
182
+ mask = _palette_masks(score, max_colors, mat)
183
+ score = torch.where(
184
+ mask.unsqueeze(1), score, torch.tensor(-1e9, device=score.device)
185
+ )
186
+ if max_swaps is not None and max_swaps < score.shape[1] - (1 if free is None else 0):
187
+ out = _viterbi_swaps(score, max_swaps, free)
188
+ else:
189
+ out = score.argmax(-1)
190
+ return out[0] if squeeze else out
191
+
192
+
193
+ def feasible(
194
+ dg: torch.Tensor, max_colors: Optional[int], max_swaps: Optional[int], free: Optional[int] = None
195
+ ) -> bool:
196
+ c, s = count_colors_swaps(dg, free)
197
+ return (max_colors is None or c <= max_colors) and (max_swaps is None or s <= max_swaps)
198
+
199
+
200
+ def neighbour_stacks(
201
+ dg: torch.Tensor,
202
+ n_mat: int,
203
+ max_colors: Optional[int],
204
+ max_swaps: Optional[int],
205
+ free: Optional[int] = None,
206
+ ) -> list:
207
+ """Feasible single-move neighbours of the stack ``dg`` [L]: recolour a
208
+ band, shift a band boundary by 1-2 layers, recolour one layer, replace one
209
+ palette material everywhere by any other material."""
210
+ L = dg.shape[0]
211
+ mat = _mat(free)
212
+ cur = dg.tolist()
213
+ seen = {tuple(cur)}
214
+ out = []
215
+
216
+ def add(cand):
217
+ t = tuple(cand)
218
+ if t in seen:
219
+ return
220
+ seen.add(t)
221
+ colors = len(set(cand) - {mat})
222
+ swaps = sum(1 for a, b in zip(cand[:-1], cand[1:]) if a != b)
223
+ if free is not None and cand[0] != free:
224
+ swaps += 1 # from the base to the first layer
225
+ if (max_colors is None or colors <= max_colors) and (
226
+ max_swaps is None or swaps <= max_swaps
227
+ ):
228
+ out.append(cand)
229
+
230
+ bands = []
231
+ start = 0
232
+ for i in range(1, L + 1):
233
+ if i == L or cur[i] != cur[start]:
234
+ bands.append((start, i - 1, cur[start]))
235
+ start = i
236
+ for a, b, c in bands:
237
+ for m in range(n_mat):
238
+ if m != c:
239
+ add(cur[:a] + [m] * (b - a + 1) + cur[b + 1:])
240
+ for (a0, b0, c0), (a1, b1, c1) in zip(bands[:-1], bands[1:]):
241
+ for d in (1, 2):
242
+ if b0 - d + 1 >= a0: # upper band grows downward
243
+ add(cur[: b0 - d + 1] + [c1] * d + cur[b0 + 1:])
244
+ if a1 + d - 1 <= b1: # lower band grows upward
245
+ add(cur[:a1] + [c0] * d + cur[a1 + d:])
246
+ # The base filament is always available to a layer, at no colour cost.
247
+ palette = sorted(set(cur) | ({mat} if mat is not None else set()))
248
+ for l in range(L):
249
+ for m in palette if max_colors is not None else range(n_mat):
250
+ if m != cur[l]:
251
+ add(cur[:l] + [m] + cur[l + 1:])
252
+ for p in sorted(set(cur)):
253
+ for m in range(n_mat):
254
+ if m != p: # replace (m outside the palette) or merge (m inside, the base included)
255
+ add([m if v == p else v for v in cur])
256
+ return [torch.tensor(c, dtype=dg.dtype, device=dg.device) for c in out]
257
+
258
+
259
+ def count_colors_swaps(dg: torch.Tensor, free: Optional[int] = None) -> tuple[int, int]:
260
+ """(materials in the stack besides the base filament, swaps) - ``free``
261
+ as in ``project_assignment``. With a base the print's colours are the
262
+ first number plus one, and the swaps include the change from the base to
263
+ the first layer."""
264
+ mats = set(torch.unique(dg).tolist())
265
+ mats.discard(_mat(free))
266
+ swaps = int((dg[1:] != dg[:-1]).sum().item())
267
+ if free is not None and int(dg[0]) != free:
268
+ swaps += 1
269
+ return len(mats), swaps
270
+
271
+
272
+ def base_material_index(background: torch.Tensor, material_colors: torch.Tensor) -> Optional[int]:
273
+ """Index of the filament the base is printed in, or None when the base
274
+ color is not one of the filaments (a custom color: it then takes a
275
+ color slot of its own). Matched by color, which covers the automatically
276
+ picked base as well as a filament chosen by hand."""
277
+ d = (material_colors.float() - background.float().view(1, -1)).abs().amax(dim=1)
278
+ i = int(torch.argmin(d))
279
+ return i if float(d[i]) < 0.5 / 255 else None