jianhey commited on
Commit
21bd12f
Β·
verified Β·
1 Parent(s): 2cd0dd5

Add model card with NVIDIA Open Model License + architecture + dataset details

Browse files
Files changed (1) hide show
  1. README.md +157 -0
README.md ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: other
3
+ license_name: nvidia-open-model-license
4
+ license_link: https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/
5
+ library_name: pytorch
6
+ pipeline_tag: image-to-image
7
+ base_model: nvidia/Cosmos-Predict2-2B-Text2Image
8
+ tags:
9
+ - anomaly-detection
10
+ - synthetic-data
11
+ - pcb
12
+ - inpainting
13
+ - cosmos
14
+ - anomalygen
15
+ ---
16
+
17
+ # Model Overview
18
+
19
+ ### Description:
20
+ Cosmos AnomalyGen β€” PCB (UC1) generates synthetic printed-circuit-board anomaly images by inpainting a user-supplied binary mask onto a clean reference PCB image, conditioned on one of three trained `<texture>+<anomaly_type>` pairs (`IC+bridge`, `passive_component+excess_solder`, `passive_component+missing`). The release ships only the few-shot-finetuned modules β€” a set of anomaly-token embeddings and a 2-layer MLP adapter β€” which plug into the frozen Cosmos-Predict2 2B Text-to-Image diffusion backbone (also using a frozen NV-DINOv2 mask encoder and a frozen T5 text encoder) at inference time. Cosmos AnomalyGen β€” UC1 v1.0.0 was developed by NVIDIA as part of the Cosmos AnomalyGen pipeline. This model is ready for commercial use.<br>
21
+
22
+ ### License/Terms of Use:
23
+
24
+ Use of the AnomalyGen finetuned modules in this release is governed by the [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/).<br>
25
+
26
+ Inference also requires the following components, which are **not** redistributed in this release and remain governed by their own terms:
27
+ - Cosmos-Predict2-2B-Text2Image β€” [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/)
28
+ - NV-DINOv2 classification model β€” distributed via NVIDIA NGC under the NVIDIA TAO license
29
+ - google-t5/t5-large text encoder β€” Apache 2.0
30
+ - The Anomaly Diffusion pipeline concept (adopted as the framework) β€” MIT License<br>
31
+
32
+ ### Deployment Geography:
33
+ Global<br>
34
+
35
+ ### Use Case:
36
+ Industrial visual-inspection teams responsible for PCB QA who have only a small number of real anomaly examples (≀62 per defect type). The model produces large-scale synthetic anomaly datasets (clean PCB + binary mask β†’ realistic bridge / excess_solder / missing-component image) for training downstream defect-detection or segmentation models, including downstream TAO toolkit consumers via the DAFT v3.0 export path. Unlike UC2 and UC3, UC1 spans two PCB texture categories (`IC` and `passive_component`), so a single checkpoint can cover defects whose appearance depends on which board region (IC area vs. passive-component area) they occur in.<br>
37
+
38
+ ### Release Date:
39
+ Github 06/02/2026 via https://github.com/NVIDIA/paidf-anomalygen<br>
40
+
41
+ ## References(s):
42
+ - Anomaly Diffusion (AAAI 2024) β€” paper: https://arxiv.org/abs/2312.05767, code: https://github.com/sjtuplayer/anomalydiffusion
43
+ - Cosmos-Predict2 β€” https://github.com/nvidia-cosmos/cosmos-predict2
44
+ - NV-DINOv2 classification model β€” https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/nv_dinov2_classification_model<br>
45
+
46
+ ## Model Architecture:
47
+ **Architecture Type:** Transformer (diffusion DiT backbone with learnable conditioning modules)<br>
48
+
49
+ **Network Architecture:**
50
+ - `anomaly_embedding` *(trainable, included in this release)*: token embeddings (256 tokens per `<texture>+<anomaly_type>` pair) β€” three pairs trained for UC1: `IC+bridge`, `passive_component+excess_solder`, `passive_component+missing`.
51
+ - `adapter` *(trainable, included in this release)*: 2-layer MLP with GELU activations (input / output hidden size = 1024), projecting the mask encoder output into the diffusion DiT conditioning space.
52
+ - `mask_encoder` *(frozen, not redistributed in this release)*: NV-DINOv2 (ViT-L) backbone with adaptive pool (kernel = 7); weights are loaded from the separately downloaded NV-DINOv2 classification checkpoint at inference time.
53
+ - `text_encoder` *(frozen, not redistributed in this release)*: google-t5/t5-large.
54
+ - These modules condition the **frozen** Cosmos-Predict2 2B T2I DiT denoiser at inference time.<br>
55
+
56
+ **This model was developed based on Cosmos-Predict2-2B-Text2Image.**<br>
57
+
58
+ **Number of model parameters:** Approximately a few million trainable parameters in the released modules (anomaly_embedding + MLP adapter), distributed as the `model/iter_000014000.pt` checkpoint file. The frozen Cosmos-Predict2 2B base contributes ~2.0Γ—10^9 (2 billion) parameters used at inference time but **not** redistributed in this release.<br>
59
+
60
+ ## Input(s):
61
+ **Input Type(s):** Image, Image (binary mask), Text<br>
62
+
63
+ **Input Format(s):**
64
+ - Image: PNG / JPG, RGB
65
+ - Mask: PNG / JPG, single-channel binary (0 = background, 255 = anomaly region; binarized at threshold 127)
66
+ - Text: anomaly-type string in the form `<texture>+<anomaly_type>` (one of `IC+bridge`, `passive_component+excess_solder`, `passive_component+missing`)<br>
67
+
68
+ **Input Parameters:**
69
+ - Image: Two-Dimensional (2D)
70
+ - Mask: Two-Dimensional (2D)
71
+ - Text: One-Dimensional (1D)<br>
72
+
73
+ **Other Properties Related to Input:** Input clean image and paired mask must have the same dimensions; the model was trained at 512Γ—512 and inference is run at the same resolution. `anomaly_type` must exactly match one of the three pairs trained for this UC1 checkpoint β€” passing an unsupported defect string is rejected by `scripts/anomaly_gen/sdg-inference/validate_jsonl.py` against this checkpoint's `ag_config.yaml β†’ dataloader_train.dataset.anomaly_types`. Because UC1 spans two textures, the chosen texture (`IC` vs. `passive_component`) must match the board region from which the clean reference image was cropped, otherwise the generated defect may look misplaced. The optional Automatic Mask Placement (AMP) tool can constrain mask placement to legal ROIs (e.g., only on IC pads, only on passive-component pads).<br>
74
+
75
+ ## Output(s)
76
+
77
+ **Output Type(s):** Image<br>
78
+
79
+ **Output Format(s):** PNG, RGB<br>
80
+
81
+ **Output Parameters:** Two-Dimensional (2D)<br>
82
+
83
+ **Other Properties Related to Output:** 512Γ—512 RGB synthetic anomaly image. Anomaly content is generated inside the user-supplied mask region; in the default `crop_and_paste=True` flow the inpainted patch is pasted back onto the clean reference image so non-masked pixels remain identical to the input. Optionally Poisson blending can be enabled. Generation metadata (per-sample guidance, crop_ratio, seed, etc.) is written to `SDG_result.csv` alongside the images.<br>
84
+
85
+ Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA's hardware (e.g. GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions.<br>
86
+
87
+ ## Software Integration:
88
+ **Runtime Engine(s):**
89
+ - PyTorch (via the Cosmos-Predict2 2B T2I pipeline)
90
+ - Cosmos AnomalyGen scripts (`scripts.anomaly_gen.synthetic_dataset_generation`, torchrun-based)
91
+ - NVIDIA TAO Toolkit β€” interop via DAFT v3.0 export (`scripts.anomaly_gen.convert_to_daft_format`)<br>
92
+
93
+ **Supported Hardware Microarchitecture Compatibility:**
94
+ - NVIDIA Ampere (A100)
95
+ - NVIDIA Hopper (H100)
96
+ - NVIDIA RTX 6000<br>
97
+
98
+ **Supported Operating System(s):**
99
+ - Linux<br>
100
+
101
+ The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific data to ensure safe and effective deployment. Following the V-model methodology, iterative testing and validation at both unit and system levels are essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment.<br>
102
+
103
+ ## Model Version(s):
104
+ v1.0.0 β€” `uc1-pcb-2b-512-iter14000` (trained 14,000 iterations; released artifact is the `model/iter_000014000.pt` checkpoint file containing finetuned modules only).<br>
105
+
106
+ ## Training, Testing, and Evaluation Datasets:
107
+
108
+ ### Dataset Overview
109
+ - Total Number of Datasets: 1 (Cosmos AnomalyGen PCB reference dataset, pre-packaged on NVIDIA NGC β€” no local preparation required)
110
+ - Total Size: 86 anomaly RGB images + 86 paired binary masks + clean PCB reference images (clean images are used at inference time, not as supervised training targets)
111
+ - Per `<texture>+<anomaly_type>` pair: `IC+bridge` = 8 anomaly images, `passive_component+excess_solder` = 16 anomaly images, `passive_component+missing` = 62 anomaly images
112
+ - Dataset partition: 100% Training (few-shot fine-tuning regime; no held-out validation or test split)<br>
113
+
114
+ ## Public Datasets
115
+ - `nv-metropolis-dev/metropolis-sdg/cosmos-anomalygen-pcb-dataset:1.0` β€” NVIDIA NGC β€” anomaly images, paired binary masks, clean reference images, and `defect_spec.jsonl`.
116
+ - `nv-metropolis-dev/metropolis-sdg/cosmos-anomalygen-pcb-assets:1.0` β€” NVIDIA NGC β€” supporting assets for the UC1 pipeline.<br>
117
+
118
+ ## Training Dataset:
119
+
120
+ ### Data Modality:
121
+ - Image<br>
122
+
123
+ ### Training Data Size:
124
+ - Less than a Million Images (86 anomaly images + 86 paired masks across three `<texture>+<anomaly_type>` pairs)
125
+
126
+ ** Data Collection Method by dataset
127
+ - Manually-Collected (camera-captured PCB images)
128
+
129
+ ** Labeling Method by dataset
130
+ - Manually-Labeled (per-defect binary masks produced for the NGC dataset)
131
+
132
+ **Properties:** 86 RGB images of printed circuit boards spanning two textures (`IC`, `passive_component`) and three defect categories β€” `IC+bridge` (8 images), `passive_component+excess_solder` (16 images), `passive_component+missing` (62 images) β€” each paired with a binary mask indicating defect pixels. Clean (no-defect) PCB images from the same NGC dataset are used as inpainting references at inference time. The class imbalance across the three pairs (e.g., 62Γ— `missing` vs. 8Γ— `bridge`) is intentional and reflects what was available; the diffusion fine-tuning is robust to per-class sample counts as low as a handful. No personal data, copyrighted human-subject content, or human subjects are present in the dataset.<br>
133
+
134
+ ### Testing Dataset:
135
+ Not Applicable β€” the model is qualitatively evaluated via the `log_image` validation callback (inpainted samples logged every N training steps) rather than against a held-out test split.<br>
136
+
137
+ ### Evaluation Dataset:
138
+ Not Applicable for this release.<br>
139
+
140
+ ## Inference:
141
+ **Acceleration Engine:** PyTorch (native, FP32). Multi-GPU rank-sharded inference is supported via `torchrun --nproc_per_node=<N>` with the `predict2_anomaly_gen_fsdp_2b` experiment.<br>
142
+
143
+ **Test Hardware:**
144
+ - NVIDIA A100
145
+ - NVIDIA H100
146
+ - NVIDIA RTX 6000<br>
147
+
148
+ ## Ethical Considerations:
149
+ NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. Developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.<br>
150
+
151
+ Please make sure you have proper rights and permissions for all input image content used as clean references or anomaly examples. PCB images are inanimate objects, but users should still verify that any incidentally captured personally identifiable content (e.g., visible serial numbers, hand-written part identifiers) is handled in accordance with applicable privacy laws prior to use.<br>
152
+
153
+ Users are responsible for model inputs and outputs. Users are responsible for ensuring safe integration of this model, including implementing guardrails as well as other safety mechanisms, prior to deployment.<br>
154
+
155
+ For more detailed information on ethical considerations for this model, please see the Bias, Explainability, Privacy, and Safety & Security subcards alongside this overview in `modelcard/UC1/`.<br>
156
+
157
+ Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).<br>