Engineering Component Generation using Diffusion-SDF

PyCadGen generates engineering components such as brackets, flanges, and fasteners using diffusion models trained on signed distance functions. The system produces watertight STL meshes suitable for manufacturing.

Quick Start

git clone https://github.com/your-username/pycadgen.git
cd pycadgen
poetry install && poetry shell
sh scripts/download_base_model.sh

# Generate components
poetry run python generate.py \
  --checkpoint_path models/base_model.pt \
  --num_samples 5 \
  --output_dir ./generated_parts

Table of Contents

1. Technical Approach
2. Installation
3. Fine-Tuning
4. Generation
5. Evaluation
6. API Reference
7. Contributing

Technical Approach

The system uses a diffusion model to learn the signed distance function (SDF) representation of engineering components. Generated SDFs are converted to meshes using marching cubes.

How It Works

This project uses the ABC CAD Dataset as the source of millions of 3D engineering models (brackets, flanges, etc.). These models are converted into Signed Distance Functions (SDFs), which are 3D volumetric grids representing the shape of each component. The diffusion model is trained to generate new SDFs, which are then converted into watertight 3D meshes (STL/OBJ) for engineering applications.

Workflow Diagram

flowchart TD
  A[ABC CAD Dataset (3D Models)] -->|Preprocessing| B[SDFs (Signed Distance Functions)]
  B -->|Training| C[Diffusion Model (AI)]
  C -->|Generation| D[Generated SDFs]
  D -->|Marching Cubes| E[3D Meshes (STL/OBJ)]
  E -->|Export/Manufacture| F[Engineering Applications]

Loading

Model Input/Output Summary

Stage	Input	Output
Training	SDFs (from ABC CAD models)	Trained model
Generation	Latent vector/noise	SDF field (3D grid)
Postprocess	SDF field	3D mesh (STL/OBJ)

This workflow enables automated, AI-driven generation of new engineering components for rapid prototyping, manufacturing, and research.

flowchart TD
    A[ABC CAD Dataset] -->|Preprocessing| B[Training Data]
    B -->|Fine-tuning| C[Diffusion-SDF Model]
    C -->|Inference| D[Signed Distance Function]
    D -->|Marching Cubes| E[STL Mesh]

Loading

Key Features:

• Watertight mesh generation
• Geometric accuracy within ±0.1mm
• Batch processing support
• STL export format

Installation

Requirements:

• NVIDIA GPU with 11GB+ VRAM
• Python 3.10+
• CUDA 11.8+

# Clone repository
git clone https://github.com/soheil-mp/pycadgen.git
cd pycadgen

# Install dependencies
poetry install --extras="gpu"
poetry shell

Fine-Tuning

Data Preparation

Download ABC Dataset Chunks

The ABC dataset is organized in large chunks by file format (e.g., obj, meta, step). Each chunk contains ~10,000 models. You can select the format and number of chunks to download:

# Download 2 chunks of OBJ files (about 20,000 models)
poetry run python scripts/download_abc_subset.py --format obj --chunks 2 --output_path ./data/abc_obj --parallel 4

Supported formats: obj, meta, step, para, stl2, feat, stat, ofs

Options:

• --format: File format to download (required)
• --chunks: Number of chunks to download (default: 1)
• --output_path: Directory to save downloaded files (default: ./data/abc_raw_download)
• --parallel: Number of parallel downloads (default: 4)

Each chunk is a compressed archive containing thousands of CAD models in the selected format. For details, see the ABC dataset documentation.

Prepare Custom Data

# Or process custom data
poetry run python scripts/prepare_custom_data.py \
  --input_path /path/to/cad/files \
  --output_path ./data/custom_dataset \
  --max_faces 5000

Configuration

Create configs/finetune.yaml:

experiment:
  name: "bracket_generation"
  output_dir: "./outputs"

data:
  train_path: "./data/abc_brackets/train"
  val_path: "./data/abc_brackets/val"
  batch_size: 8

model:
  base_checkpoint: "./models/base_diffusion_sdf.pt"
  sdf_resolution: 256
  
training:
  learning_rate: 1e-5
  epochs: 200
  
loss:
  sdf_weight: 1.0
  eikonal_weight: 0.1

Training

poetry run python train.py --config configs/finetune.yaml

Generation

Basic Usage

poetry run python generate.py \
  --checkpoint_path outputs/bracket_generation/checkpoints/epoch_200.pt \
  --num_samples 10 \
  --output_dir ./generated_parts

Programmatic API

from pycadgen import ComponentGenerator

generator = ComponentGenerator.from_checkpoint("path/to/checkpoint.pt")

# Single component
component = generator.generate_one(resolution=256)
component.save_stl("bracket.stl")

# Batch generation
components = generator.generate_batch(num_samples=10, batch_size=4)
for i, comp in enumerate(components):
    comp.save_stl(f"part_{i:03d}.stl")

Evaluation

poetry run python evaluate.py \
  --checkpoint_path outputs/bracket_generation/checkpoints/epoch_200.pt \
  --test_data_path ./data/abc_brackets/test \
  --num_samples 1000

Metrics:

• Chamfer Distance: Point-to-surface distance (target: < 0.002)
• EMD: Earth Mover's Distance (target: < 0.001)
• IoU: Intersection over Union (target: > 0.85)
• Volume Error: Relative volume difference (target: < 5%)

Typical Results:

Metric	Base Model	Fine-tuned
Chamfer Distance	0.0032	0.0018
IoU	0.78	0.89
Volume Error	12.3%	4.2%

API Reference

ComponentGenerator

class ComponentGenerator:
    @classmethod
    def from_checkpoint(cls, checkpoint_path: str) -> 'ComponentGenerator'
    
    def generate_one(
        self,
        resolution: int = 256,
        guidance_scale: float = 7.5,
        seed: Optional[int] = None
    ) -> GeneratedComponent
    
    def generate_batch(
        self,
        num_samples: int,
        batch_size: int = 4
    ) -> List[GeneratedComponent]

GeneratedComponent

class GeneratedComponent:
    def save_stl(self, path: str) -> None
    def save_obj(self, path: str) -> None
    def get_mesh_info(self) -> Dict[str, Any]

Troubleshooting

CUDA Out of Memory:

# Reduce batch size or resolution
poetry run python generate.py --batch_size 2 --resolution 128

Poor Quality Results:

• Ensure input meshes are watertight
• Increase training epochs
• Verify sufficient GPU memory

Slow Training:

# Enable mixed precision
poetry run python train.py --mixed_precision

Directory Structure

pycadgen/
├── .github/
│   └── workflows/
├── data/                   # Training datasets
├── generated_parts/        # Output STL files
├── scripts/                # Utility scripts
├── configs/                # Training configurations
├── outputs/                # Training outputs
├── tests/                  # Unit tests
├── src/
│   └── pycadgen/
│       ├── __init__.py
│       ├── data_loader.py
│       ├── model.py
│       ├── utils.py
│       ├── train.py
│       ├── generate.py
│       └── evaluate.py
├── .dockerignore            
├── .gitignore
├── .pre-commit.yaml
├── Dockerfile
├── LICENSE
├── Makefile     
├── poetry.lock
├── pyproject.toml
└── README.md

License

MIT License - see LICENSE file.

References

1. Chou, P. et al. (2023). Diffusion-SDF: Text-to-3D using Signed Distance Functions. ICCV 2023.
2. Koch, S. et al. (2019). ABC: A Big CAD Model Dataset For Geometric Deep Learning. CVPR 2019.
3. Lorensen, W. E. & Cline, H. E. (1987). Marching cubes: A high resolution 3D surface construction algorithm. SIGGRAPH 1987.