autoforge 1.8.0__tar.gz → 1.8.2__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 (41) hide show
  1. {autoforge-1.8.0 → autoforge-1.8.2}/PKG-INFO +1 -18
  2. {autoforge-1.8.0 → autoforge-1.8.2}/README.md +0 -17
  3. {autoforge-1.8.0 → autoforge-1.8.2}/pyproject.toml +1 -1
  4. {autoforge-1.8.0 → autoforge-1.8.2}/src/AutoForge.egg-info/PKG-INFO +1 -18
  5. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/OptimizerHelper.py +0 -35
  6. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/PruningHelper.py +121 -81
  7. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Modules/Optimizer.py +57 -17
  8. {autoforge-1.8.0 → autoforge-1.8.2}/setup.cfg +0 -0
  9. {autoforge-1.8.0 → autoforge-1.8.2}/src/AutoForge.egg-info/SOURCES.txt +0 -0
  10. {autoforge-1.8.0 → autoforge-1.8.2}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  11. {autoforge-1.8.0 → autoforge-1.8.2}/src/AutoForge.egg-info/entry_points.txt +0 -0
  12. {autoforge-1.8.0 → autoforge-1.8.2}/src/AutoForge.egg-info/requires.txt +0 -0
  13. {autoforge-1.8.0 → autoforge-1.8.2}/src/AutoForge.egg-info/top_level.txt +0 -0
  14. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Gui/__init__.py +0 -0
  15. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Gui/browser_gui/Views/__init__.py +0 -0
  16. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Gui/browser_gui/Views/backend_database.py +0 -0
  17. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Gui/browser_gui/Views/filament_editor.py +0 -0
  18. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Gui/browser_gui/app.py +0 -0
  19. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Gui/browser_gui/main.py +0 -0
  20. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/CAdamW.py +0 -0
  21. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/FilamentHelper.py +0 -0
  22. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/Heightmaps/ChristofidesHeightMap.py +0 -0
  23. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/Heightmaps/DepthEstimateHeightMap.py +0 -0
  24. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/Heightmaps/OldHeightmap.py +0 -0
  25. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/Heightmaps/__init__.py +0 -0
  26. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/ImageHelper.py +0 -0
  27. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/OutputHelper.py +0 -0
  28. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/__init__.py +0 -0
  29. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Helper/filamentcolors_library.py +0 -0
  30. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Init_comparer/__init__.py +0 -0
  31. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Init_comparer/auto_forge_init_check.py +0 -0
  32. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Init_comparer/plot_results.py +0 -0
  33. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Loss/LossFunctions.py +0 -0
  34. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  35. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Loss/__init__.py +0 -0
  36. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Misc/__init__.py +0 -0
  37. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Misc/color_fit_all.py +0 -0
  38. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/Modules/__init__.py +0 -0
  39. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/__init__.py +0 -0
  40. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/__main__.py +0 -0
  41. {autoforge-1.8.0 → autoforge-1.8.2}/src/autoforge/auto_forge.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: AutoForge
3
- Version: 1.8.0
3
+ Version: 1.8.2
4
4
  Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
5
5
  Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
6
  License: CC BY-NC-SA 4.0
@@ -66,22 +66,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
66
66
  - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
67
67
  - **Hueforge export**: Outputs a project file that can be opened with hueforge.
68
68
 
69
- ## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
70
-
71
- If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
72
-
73
- 1. **Open the Notebook:**
74
- Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
75
-
76
- 2. **Select GPU Runtime:**
77
- In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
78
-
79
- 3. **Run All Cells:**
80
- Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
81
-
82
- 4. **Review and Download:**
83
- After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
84
-
85
69
  ## Manual Installation
86
70
 
87
71
  To install AutoForge, simply install the current version from PyPI:
@@ -116,7 +100,6 @@ autoforge --input_image path/to/input_image.jpg --json_file %APPDATA%\HueForge\F
116
100
  autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
117
101
  ```
118
102
 
119
-
120
103
  ### Command Line Arguments
121
104
 
122
105
  - `--config` *(Optional)* Path to a configuration file with the settings.
@@ -40,22 +40,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
40
40
  - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
41
41
  - **Hueforge export**: Outputs a project file that can be opened with hueforge.
42
42
 
43
- ## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
44
-
45
- If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
46
-
47
- 1. **Open the Notebook:**
48
- Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
49
-
50
- 2. **Select GPU Runtime:**
51
- In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
52
-
53
- 3. **Run All Cells:**
54
- Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
55
-
56
- 4. **Review and Download:**
57
- After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
58
-
59
43
  ## Manual Installation
60
44
 
61
45
  To install AutoForge, simply install the current version from PyPI:
@@ -90,7 +74,6 @@ autoforge --input_image path/to/input_image.jpg --json_file %APPDATA%\HueForge\F
90
74
  autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
91
75
  ```
92
76
 
93
-
94
77
  ### Command Line Arguments
95
78
 
96
79
  - `--config` *(Optional)* Path to a configuration file with the settings.
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="1.8.0"
7
+ version="1.8.2"
8
8
  description = "AutoForge is a Python tool for generating 3D printed layered models from an input image."
9
9
  authors = [{ name = "Hendric Voss", email = "hvoss@techfak.uni-bielefeld.de" }]
10
10
  readme = "README.md"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: AutoForge
3
- Version: 1.8.0
3
+ Version: 1.8.2
4
4
  Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
5
5
  Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
6
  License: CC BY-NC-SA 4.0
@@ -66,22 +66,6 @@ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at m
66
66
  - **Live Visualization**: (Optional) Displays live composite images during the optimization process.
67
67
  - **Hueforge export**: Outputs a project file that can be opened with hueforge.
68
68
 
69
- ## [Google Colab Notebook](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih)
70
-
71
- If you're new to Google Colab, follow these simple steps to run autoforge and get your output:
72
-
73
- 1. **Open the Notebook:**
74
- Click [here](https://colab.research.google.com/drive/13sWCss9-ew2xnoJl4c4SSpComSNi-1ih) to open the notebook in Google Colab.
75
-
76
- 2. **Select GPU Runtime:**
77
- In Colab, go to the menu and click **Runtime → Change runtime type**. Under **Hardware accelerator**, choose **T4 GPU** (or any other available GPU type).
78
-
79
- 3. **Run All Cells:**
80
- Once the notebook loads, click **Runtime → Run all** to execute the cells in sequence. Follow the on-screen prompts to configure the materials CSV, upload your image, and run Autoforge.
81
-
82
- 4. **Review and Download:**
83
- After the process completes, you'll see a preview of the output image. Finally, the notebook will zip the output folder for you to download.
84
-
85
69
  ## Manual Installation
86
70
 
87
71
  To install AutoForge, simply install the current version from PyPI:
@@ -116,7 +100,6 @@ autoforge --input_image path/to/input_image.jpg --json_file %APPDATA%\HueForge\F
116
100
  autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv --pruning_max_colors 8 --pruning_max_swaps 20
117
101
  ```
118
102
 
119
-
120
103
  ### Command Line Arguments
121
104
 
122
105
  - `--config` *(Optional)* Path to a configuration file with the settings.
@@ -277,38 +277,3 @@ def composite_image_disc(
277
277
  comp = comp + rem_after.unsqueeze(-1) * background # [H,W,3]
278
278
 
279
279
  return comp * 255.0
280
-
281
-
282
- def discretize_solution(
283
- params: dict, tau_global: float, h: float, max_layers: int, rng_seed: int = -1
284
- ):
285
- """
286
- Convert continuous logs to discrete layer counts and discrete color IDs.
287
-
288
- Args:
289
- params (dict): Dictionary containing the parameters 'pixel_height_logits' and 'global_logits'.
290
- tau_global (float): Temperature parameter for global material assignment.
291
- h (float): Height of each layer.
292
- max_layers (int): Maximum number of layers.
293
- rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
294
-
295
- Returns:
296
- tuple: A tuple containing:
297
- - torch.Tensor: Discrete global material assignments, shape [max_layers].
298
- - torch.Tensor: Discrete height image, shape [H, W].
299
- """
300
- pixel_height_logits = params["pixel_height_logits"]
301
- global_logits = params["global_logits"]
302
- pixel_heights = (max_layers * h) * torch.sigmoid(pixel_height_logits)
303
- discrete_height_image = torch.round(pixel_heights / h).to(torch.int32)
304
- discrete_height_image = torch.clamp(discrete_height_image, 0, max_layers)
305
-
306
- num_layers = global_logits.shape[0]
307
- discrete_global_vals = []
308
- for j in range(num_layers):
309
- p = deterministic_gumbel_softmax(
310
- global_logits[j], tau_global, hard=True, rng_seed=rng_seed + j
311
- )
312
- discrete_global_vals.append(torch.argmax(p))
313
- discrete_global = torch.stack(discrete_global_vals, dim=0)
314
- return discrete_global, discrete_height_image
@@ -6,7 +6,7 @@ from tqdm import tqdm
6
6
  from joblib import Parallel, delayed
7
7
  import threading
8
8
 
9
- from autoforge.Helper.OptimizerHelper import discretize_solution, composite_image_disc
9
+ from autoforge.Helper.OptimizerHelper import composite_image_disc
10
10
  from autoforge.Loss.LossFunctions import compute_loss
11
11
  from autoforge.Modules.Optimizer import FilamentOptimizer
12
12
 
@@ -325,7 +325,13 @@ def merge_bands(
325
325
 
326
326
 
327
327
  def remove_layer_from_solution(
328
- params, layer_to_remove, final_tau, h, current_max_layers, rng_seed
328
+ optimizer,
329
+ params,
330
+ layer_to_remove,
331
+ final_tau,
332
+ h,
333
+ current_max_layers,
334
+ rng_seed,
329
335
  ):
330
336
  """
331
337
  Remove one layer from the solution.
@@ -343,8 +349,15 @@ def remove_layer_from_solution(
343
349
  new_max_layers (int): Updated number of layers.
344
350
  """
345
351
  # Get the current discrete height image (used to decide which pixels need adjusting)
346
- _, disc_height = discretize_solution(
347
- params, final_tau, h, current_max_layers, rng_seed
352
+ effective_logits = optimizer._apply_height_offset(
353
+ params["pixel_height_logits"], params["height_offsets"]
354
+ )
355
+ _, disc_height = optimizer.discretize_solution(
356
+ params,
357
+ final_tau,
358
+ h,
359
+ current_max_layers,
360
+ rng_seed,
348
361
  )
349
362
 
350
363
  # Remove the candidate layer from the global (color) assignment.
@@ -358,25 +371,24 @@ def remove_layer_from_solution(
358
371
  new_max_layers = new_global_logits.shape[0]
359
372
 
360
373
  # Compute current effective height: height = (current_max_layers * h) * sigmoid(pixel_height_logits)
361
- current_height = (
362
- current_max_layers * h * torch.sigmoid(params["pixel_height_logits"])
363
- )
364
-
365
- # For pixels where the discrete height is at or above the removed layer, subtract h.
374
+ current_height = current_max_layers * h * torch.sigmoid(effective_logits)
366
375
  new_height = current_height.clone()
367
- mask = disc_height >= layer_to_remove # <-- Changed from > to >= here
376
+ mask = disc_height >= layer_to_remove
368
377
  new_height[mask] = new_height[mask] - h
369
378
 
370
379
  # Invert the sigmoid mapping:
371
380
  # We need new_pixel_height_logits such that:
372
381
  # sigmoid(new_pixel_height_logits) = new_height / (new_max_layers * h)
373
382
  eps = 1e-6
374
- new_ratio = torch.clamp(new_height / (new_max_layers * h), eps, 1 - eps)
375
- new_pixel_height_logits = torch.log(new_ratio) - torch.log(1 - new_ratio)
376
-
383
+ ratio = torch.clamp(new_height / (new_max_layers * h), eps, 1.0 - eps)
384
+ new_effective_logits = torch.log(ratio) - torch.log1p(-ratio)
385
+ new_pixel_height_logits = new_effective_logits - (
386
+ effective_logits - params["pixel_height_logits"]
387
+ )
377
388
  new_params = {
378
389
  "global_logits": new_global_logits,
379
390
  "pixel_height_logits": new_pixel_height_logits,
391
+ "height_offsets": params["height_offsets"],
380
392
  }
381
393
  return new_params, new_max_layers
382
394
 
@@ -399,129 +411,157 @@ def prune_redundant_layers(
399
411
  stopping as soon as one batch yields an improvement.
400
412
  """
401
413
  current_max_layers = optimizer.best_params["global_logits"].shape[0]
402
- optimizer.max_layers = current_max_layers
403
-
404
- def score_layer(
405
- layer_idx: int,
406
- base_params: dict,
407
- final_tau: float,
408
- h: float,
409
- cur_max_layers: int,
410
- best_seed: int,
411
- ):
414
+ optimizer.max_layers = current_max_layers # keep optimiser in sync
415
+
416
+ # Baseline loss with current best parameters
417
+ with _gpu_lock, torch.no_grad():
418
+ eff_logits = optimizer._apply_height_offset(
419
+ optimizer.best_params["pixel_height_logits"],
420
+ optimizer.best_params["height_offsets"],
421
+ )
422
+ ref_comp = composite_image_disc(
423
+ eff_logits,
424
+ optimizer.best_params["global_logits"],
425
+ optimizer.vis_tau,
426
+ optimizer.vis_tau,
427
+ optimizer.h,
428
+ current_max_layers,
429
+ optimizer.material_colors,
430
+ optimizer.material_TDs,
431
+ optimizer.background,
432
+ rng_seed=optimizer.best_seed,
433
+ )
434
+ best_loss = compute_loss(ref_comp, optimizer.target).item()
435
+
436
+ tbar = tqdm(
437
+ desc=f"Layer pruning | Loss {best_loss:.4f}",
438
+ total=max(current_max_layers - 1, 0),
439
+ leave=False,
440
+ )
441
+ removed_layers = 0
442
+ improvement = True
443
+
444
+ # ----------------------------------------------------------
445
+ # Inner helper: evaluate the effect of removing layer idx
446
+ # ----------------------------------------------------------
447
+ def score_layer(idx: int) -> tuple[float, dict, int]:
412
448
  cand_params, cand_max_layers = remove_layer_from_solution(
413
- base_params, layer_idx, final_tau, h, cur_max_layers, best_seed
449
+ optimizer,
450
+ optimizer.best_params,
451
+ idx,
452
+ optimizer.vis_tau,
453
+ optimizer.h,
454
+ current_max_layers,
455
+ optimizer.best_seed,
456
+ )
457
+ eff_logits = optimizer._apply_height_offset(
458
+ cand_params["pixel_height_logits"], cand_params["height_offsets"]
414
459
  )
415
460
  with _gpu_lock, torch.no_grad():
416
461
  cand_comp = composite_image_disc(
417
- cand_params["pixel_height_logits"],
462
+ eff_logits,
418
463
  cand_params["global_logits"],
419
- final_tau,
420
- final_tau,
421
- h,
464
+ optimizer.vis_tau,
465
+ optimizer.vis_tau,
466
+ optimizer.h,
422
467
  cand_max_layers,
423
468
  optimizer.material_colors,
424
469
  optimizer.material_TDs,
425
470
  optimizer.background,
426
- rng_seed=best_seed,
471
+ rng_seed=optimizer.best_seed,
427
472
  )
428
473
  cand_loss = compute_loss(cand_comp, optimizer.target).item()
429
474
  return cand_loss, cand_params, cand_max_layers
430
475
 
431
- with _gpu_lock, torch.no_grad():
432
- comp = optimizer.get_best_discretized_image()
433
- best_loss = compute_loss(comp, optimizer.target).item()
434
-
435
- tbar = tqdm(
436
- desc=f"Layer pruning | Loss {best_loss:.4f}",
437
- total=max(current_max_layers - 1, 0),
438
- leave=False,
439
- )
440
- removed_layers = 0
441
- improvement = True
442
-
476
+ # ----------------------------------------------------------
477
+ # Main pruning loop
478
+ # ----------------------------------------------------------
443
479
  while current_max_layers > pruning_min_layers and (
444
480
  improvement or current_max_layers > pruning_max_layers
445
481
  ):
446
482
  tbar.update(1)
447
483
  improvement = False
448
-
449
484
  layer_indices = list(range(current_max_layers))
485
+
450
486
  if fast:
451
487
  random.shuffle(layer_indices)
452
488
  chunk_size = max(1, math.ceil(len(layer_indices) * chunking_percent))
453
- c_cand = None
454
- c_loss = 1000
489
+ best_candidate = None
490
+ best_cand_loss = float("inf")
491
+
455
492
  for chunk in _chunked(layer_indices, chunk_size):
456
493
  cand_results = Parallel(
457
494
  n_jobs=n_jobs, backend="threading", prefer="threads"
458
- )(
459
- delayed(score_layer)(
460
- idx,
461
- optimizer.best_params,
462
- optimizer.vis_tau,
463
- optimizer.h,
464
- current_max_layers,
465
- optimizer.best_seed,
466
- )
467
- for idx in chunk
468
- )
495
+ )(delayed(score_layer)(idx) for idx in chunk)
496
+
469
497
  cand_loss, cand_params, cand_max_layers = min(
470
498
  cand_results, key=lambda x: x[0]
471
499
  )
500
+
472
501
  if cand_loss <= best_loss * (1 + allowed_loss_increase_percent):
502
+ # Accept immediately and restart outer loop
473
503
  removed_layers += 1
474
504
  best_loss = cand_loss
475
505
  current_max_layers = cand_max_layers
476
506
  optimizer.best_params = cand_params
477
507
  optimizer.max_layers = current_max_layers
478
508
  tbar.set_description(
479
- f"Layer pruning | Loss {best_loss:.4f} | Removed {removed_layers} | Layers {current_max_layers}"
509
+ f"Layer pruning | Loss {best_loss:.4f} | "
510
+ f"Removed {removed_layers} | Layers {current_max_layers}"
480
511
  )
481
512
  improvement = True
482
- break # restart outer while loop with updated solution
483
- else:
484
- if current_max_layers > pruning_max_layers and cand_loss < c_loss:
485
- c_cand = (cand_params, current_max_layers)
486
- if c_cand is not None:
513
+ break
514
+
515
+ # Track the best candidate in case we must enforce a hard limit
516
+ if cand_loss < best_cand_loss:
517
+ best_cand_loss = cand_loss
518
+ best_candidate = (cand_params, cand_max_layers)
519
+
520
+ # Forced removal to meet pruning_max_layers
521
+ if (
522
+ not improvement
523
+ and current_max_layers > pruning_max_layers
524
+ and best_candidate is not None
525
+ ):
526
+ cand_params, cand_max_layers = best_candidate
487
527
  removed_layers += 1
488
- best_loss = c_loss
489
- current_max_layers = c_cand[1]
490
- optimizer.best_params = c_cand[0]
528
+ best_loss = best_cand_loss
529
+ current_max_layers = cand_max_layers
530
+ optimizer.best_params = cand_params
491
531
  optimizer.max_layers = current_max_layers
492
532
  tbar.set_description(
493
- f"Layer pruning | Loss {best_loss:.4f} | Removed {removed_layers} | Layers {current_max_layers}"
533
+ f"Layer pruning | Loss {best_loss:.4f} | "
534
+ f"Removed {removed_layers} | Layers {current_max_layers}"
494
535
  )
495
536
  improvement = True
496
537
 
497
538
  if not improvement:
498
- break
539
+ break # no chunk helped and layer count is within budget
499
540
  else:
541
+ # Exhaustive search
500
542
  cand_results = Parallel(
501
543
  n_jobs=n_jobs, backend="threading", prefer="threads"
502
- )(
503
- delayed(score_layer)(
504
- idx,
505
- optimizer.best_params,
506
- optimizer.vis_tau,
507
- optimizer.h,
508
- current_max_layers,
509
- optimizer.best_seed,
510
- )
511
- for idx in layer_indices
512
- )
544
+ )(delayed(score_layer)(idx) for idx in layer_indices)
545
+
513
546
  cand_loss, cand_params, cand_max_layers = min(
514
547
  cand_results, key=lambda x: x[0]
515
548
  )
516
- if cand_loss <= best_loss or current_max_layers > pruning_max_layers:
549
+
550
+ tbar.set_description(
551
+ f"Layer pruning | Current Candidate Loss: {cand_loss:.4f} | Best Loss {best_loss:.4f} | "
552
+ f"Removed {removed_layers} | Layers {current_max_layers}"
553
+ )
554
+
555
+ if (
556
+ cand_loss <= best_loss * (1 + allowed_loss_increase_percent)
557
+ or current_max_layers > pruning_max_layers
558
+ ):
517
559
  removed_layers += 1
518
560
  best_loss = cand_loss
519
561
  current_max_layers = cand_max_layers
520
562
  optimizer.best_params = cand_params
521
563
  optimizer.max_layers = current_max_layers
522
- tbar.set_description(
523
- f"Layer pruning | Loss {best_loss:.4f} | Removed {removed_layers} | Layers {current_max_layers}"
524
- )
564
+
525
565
  improvement = True
526
566
  else:
527
567
  break
@@ -9,9 +9,9 @@ from tqdm import tqdm
9
9
 
10
10
  from autoforge.Helper.CAdamW import CAdamW
11
11
  from autoforge.Helper.OptimizerHelper import (
12
- discretize_solution,
13
12
  composite_image_cont,
14
13
  composite_image_disc,
14
+ deterministic_gumbel_softmax,
15
15
  )
16
16
  from autoforge.Loss.LossFunctions import loss_fn, compute_loss
17
17
 
@@ -115,6 +115,7 @@ class FilamentOptimizer:
115
115
  self.params = {
116
116
  "pixel_height_logits": self.pixel_height_logits,
117
117
  "global_logits": global_logits_init,
118
+ "height_offsets": self.height_offsets,
118
119
  }
119
120
 
120
121
  # Tau schedule
@@ -303,6 +304,50 @@ class FilamentOptimizer:
303
304
 
304
305
  return loss
305
306
 
307
+ def discretize_solution(
308
+ self,
309
+ params: dict,
310
+ tau_global: float,
311
+ h: float,
312
+ max_layers: int,
313
+ rng_seed: int = -1,
314
+ ):
315
+ """
316
+ Convert continuous logs to discrete layer counts and discrete color IDs.
317
+
318
+ Args:
319
+ params (dict): Dictionary containing the parameters 'pixel_height_logits' and 'global_logits'.
320
+ tau_global (float): Temperature parameter for global material assignment.
321
+ h (float): Height of each layer.
322
+ max_layers (int): Maximum number of layers.
323
+ rng_seed (int, optional): Random seed for deterministic sampling. Defaults to -1.
324
+
325
+ Returns:
326
+ tuple: A tuple containing:
327
+ - torch.Tensor: Discrete global material assignments, shape [max_layers].
328
+ - torch.Tensor: Discrete height image, shape [H, W].
329
+ """
330
+ pixel_logits = params["pixel_height_logits"]
331
+ pixel_offset = params["height_offsets"]
332
+ effective_logits = self._apply_height_offset(
333
+ pixel_logits=pixel_logits, height_offsets=pixel_offset
334
+ )
335
+
336
+ global_logits = params["global_logits"]
337
+ pixel_heights = (max_layers * h) * torch.sigmoid(effective_logits)
338
+ discrete_height_image = torch.round(pixel_heights / h).to(torch.int32)
339
+ discrete_height_image = torch.clamp(discrete_height_image, 0, max_layers)
340
+
341
+ num_layers = global_logits.shape[0]
342
+ discrete_global_vals = []
343
+ for j in range(num_layers):
344
+ p = deterministic_gumbel_softmax(
345
+ global_logits[j], tau_global, hard=True, rng_seed=rng_seed + j
346
+ )
347
+ discrete_global_vals.append(torch.argmax(p))
348
+ discrete_global = torch.stack(discrete_global_vals, dim=0)
349
+ return discrete_global, discrete_height_image
350
+
306
351
  def log_to_tensorboard(
307
352
  self, interval: int = 100, namespace: str = "", step: int = None
308
353
  ):
@@ -489,13 +534,8 @@ class FilamentOptimizer:
489
534
  )
490
535
 
491
536
  if best:
492
- effective_logits = self._apply_height_offset(
493
- self.best_params["pixel_height_logits"],
494
- self.best_params["height_offsets"],
495
- )
496
- current_params["pixel_height_logits"] = effective_logits
497
- disc_global, disc_height_image = discretize_solution(
498
- current_params,
537
+ disc_global, disc_height_image = self.discretize_solution(
538
+ self.best_params,
499
539
  self.vis_tau,
500
540
  self.h,
501
541
  self.max_layers,
@@ -505,7 +545,7 @@ class FilamentOptimizer:
505
545
  else:
506
546
  tau_height, tau_global = self._get_tau()
507
547
  with torch.no_grad():
508
- disc_global, disc_height_image = discretize_solution(
548
+ disc_global, disc_height_image = self.discretize_solution(
509
549
  current_params,
510
550
  tau_height,
511
551
  self.h,
@@ -603,18 +643,14 @@ class FilamentOptimizer:
603
643
  tau_h, tau_g = self.vis_tau, self.vis_tau
604
644
  with torch.no_grad():
605
645
  effective_logits = self._apply_height_offset()
606
- params_with_offset = {
607
- "pixel_height_logits": effective_logits,
608
- "global_logits": self.params["global_logits"],
609
- }
610
- disc_global, disc_height_image = discretize_solution(
611
- params_with_offset, tau_g, self.h, self.max_layers, rng_seed=seed
646
+ disc_global, disc_height_image = self.discretize_solution(
647
+ self.params, tau_g, self.h, self.max_layers, rng_seed=seed
612
648
  )
613
649
 
614
650
  # 2) Compute discrete-mode composite
615
651
  with torch.no_grad():
616
652
  comp_disc = composite_image_disc(
617
- self.params["pixel_height_logits"],
653
+ effective_logits,
618
654
  self.params["global_logits"],
619
655
  self.vis_tau,
620
656
  self.vis_tau,
@@ -659,8 +695,12 @@ class FilamentOptimizer:
659
695
  best_loss = start_loss
660
696
  for i in tqdm(range(num_seeds), desc="Searching for new best seed"):
661
697
  seed = np.random.randint(0, 1000000)
662
- comp_disc = composite_image_disc(
698
+ effective_logits = self._apply_height_offset(
663
699
  self.best_params["pixel_height_logits"],
700
+ self.best_params["height_offsets"],
701
+ )
702
+ comp_disc = composite_image_disc(
703
+ effective_logits,
664
704
  self.best_params["global_logits"],
665
705
  self.vis_tau,
666
706
  self.vis_tau,
File without changes