pytorch-precision-converter 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.
@@ -0,0 +1,9 @@
1
+ # MIT License
2
+
3
+ Copyright (c) 2023 Angelo Lamonaca
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
6
+
7
+ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
8
+
9
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
@@ -0,0 +1,191 @@
1
+ Metadata-Version: 2.4
2
+ Name: pytorch-precision-converter
3
+ Version: 2.0.0
4
+ Summary: Convert PyTorch checkpoints and safetensors files to fp16, bf16 or fp8 without breaking them.
5
+ Author: Angelo Lamonaca
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/angelolamonaca/PyTorch-Precision-Converter
8
+ Project-URL: Issues, https://github.com/angelolamonaca/PyTorch-Precision-Converter/issues
9
+ Project-URL: Changelog, https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/CHANGELOG.md
10
+ Keywords: pytorch,safetensors,fp16,bf16,fp8,checkpoint,quantization,stable-diffusion,ema
11
+ Classifier: Development Status :: 5 - Production/Stable
12
+ Classifier: Environment :: Console
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
17
+ Requires-Python: >=3.9
18
+ Description-Content-Type: text/markdown
19
+ License-File: LICENSE
20
+ Requires-Dist: torch>=2.1
21
+ Requires-Dist: safetensors>=0.4.3
22
+ Requires-Dist: numpy
23
+ Provides-Extra: test
24
+ Requires-Dist: pytest>=7; extra == "test"
25
+ Dynamic: license-file
26
+
27
+ # PyTorch Precision Converter
28
+
29
+ [![PyPI](https://img.shields.io/pypi/v/pytorch-precision-converter)](https://pypi.org/project/pytorch-precision-converter/)
30
+ [![Tests](https://github.com/angelolamonaca/PyTorch-Precision-Converter/actions/workflows/ci.yml/badge.svg)](https://github.com/angelolamonaca/PyTorch-Precision-Converter/actions/workflows/ci.yml)
31
+ ![Python](https://img.shields.io/badge/python-3.9%2B-blue)
32
+ ![PyTorch](https://img.shields.io/badge/PyTorch-2.1%2B-orange)
33
+ [![License: MIT](https://img.shields.io/badge/license-MIT-green)](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE)
34
+
35
+ Make PyTorch checkpoints and `safetensors` files **half the size (a quarter with fp8)** by converting them to
36
+ `fp16`, `bf16` or `fp8`, and strip the EMA weights you don't need for inference. One command, no model code
37
+ needed, runs on the CPU.
38
+
39
+ ```console
40
+ $ precision-convert model.ckpt -p fp16 -t ema-only -st
41
+ Converted 280 tensors to fp16, left 1 unchanged, dropped 242 EMA entries
42
+ Used the EMA weights for 240 tensors
43
+ Saved model-ema-only-fp16.safetensors (29.93 MB, 27% of 109.23 MB)
44
+ ```
45
+
46
+ ## Why not just call `.half()`?
47
+
48
+ A one-line cast quietly breaks things. This tool handles the details:
49
+
50
+ - **Only floating-point tensors are cast.** Integer and bool tensors like `position_ids` or step counters keep
51
+ their dtype instead of turning into floats.
52
+ - **No silent `inf` or `NaN`.** fp16 tops out at 65,504 and fp8 at 448. Values beyond the range are clamped
53
+ to it and you get a warning, rather than infinities appearing in your weights.
54
+ - **Real EMA extraction.** `-t ema-only` puts the EMA weights where the model loads them from, for Stable
55
+ Diffusion and other latent-diffusion checkpoints (`model_ema.*`), then drops the duplicates.
56
+ - **safetensors metadata survives.** LoRA training info (`ss_*` keys) and other metadata is copied to the output.
57
+ - **Safe by default.** PyTorch checkpoints load with `weights_only=True`, so a malicious `.ckpt` can't run code.
58
+ Files that need full unpickling require an explicit `--unsafe-load`.
59
+ - **Your files are safe too.** Output is written to a temporary file and moved into place when complete, the
60
+ input is never replaced, and existing files are only overwritten with `--overwrite`.
61
+ - **Mixed precision when you need it.** `--keep-original "first_stage_model.*"` keeps matching tensors (here the
62
+ VAE) in their original dtype.
63
+
64
+ ## Install
65
+
66
+ ```bash
67
+ pip install pytorch-precision-converter
68
+ ```
69
+
70
+ For the latest unreleased changes, install from GitHub instead:
71
+ `pip install git+https://github.com/angelolamonaca/PyTorch-Precision-Converter`.
72
+
73
+ Either way you get the `precision-convert` command (`python -m precision_converter` works too). It needs Python 3.9+,
74
+ PyTorch 2.1+ and `safetensors`.
75
+
76
+ ## Usage
77
+
78
+ ```bash
79
+ precision-convert INPUT [-p PRECISION] [-t TYPE] [-st] [-o OUTPUT] [options]
80
+ ```
81
+
82
+ | Option | What it does |
83
+ |---|---|
84
+ | `INPUT` | `.safetensors` file or PyTorch checkpoint (`.ckpt`, `.pt`, `.pth`, `.bin`). `-f INPUT` works too. |
85
+ | `-p`, `--precision` | `fp32`, `fp16`, `bf16`, `fp8_e4m3fn`, `fp8_e5m2`, or `keep` (default: change no dtypes). |
86
+ | `-t`, `--type` | `full` keeps everything (default). `ema-only` uses the EMA weights. `no-ema` drops the EMA weights. |
87
+ | `-st`, `--safe-tensors` | Save as `.safetensors` instead of a PyTorch `.ckpt`. |
88
+ | `-o`, `--output` | Output path; its extension picks the format. Default: `<input>-<type>-<precision>` next to the input. |
89
+ | `--keep-original PATTERN` | Leave tensors whose key matches the glob in their original dtype. Repeatable. |
90
+ | `--info` | Show the tensors, dtypes, sizes, EMA and metadata in the file, then exit. |
91
+ | `--unsafe-load` | Fully unpickle PyTorch checkpoints. Only for files you trust. |
92
+ | `--overwrite` | Replace the output file if it exists. |
93
+ | `-v`, `--verbose` | List each tensor that had values clamped. |
94
+
95
+ ## Recipes
96
+
97
+ **Shrink a Stable Diffusion training checkpoint for inference** (EMA weights, fp16, safetensors):
98
+
99
+ ```bash
100
+ precision-convert model.ckpt -p fp16 -t ema-only -st
101
+ ```
102
+
103
+ **Convert a LoRA to fp16 and keep its training metadata:**
104
+
105
+ ```bash
106
+ precision-convert my-lora.safetensors -p fp16 -o my-lora-fp16.safetensors
107
+ ```
108
+
109
+ **Store a large model in fp8** for UIs that load fp8 weights, such as ComfyUI, keeping normalization layers
110
+ precise:
111
+
112
+ ```bash
113
+ precision-convert model.safetensors -p fp8_e4m3fn -st --keep-original "*norm*"
114
+ ```
115
+
116
+ **Keep the VAE in fp32** (fp16 VAEs can produce NaNs and black images in some models):
117
+
118
+ ```bash
119
+ precision-convert model.safetensors -p fp16 -st --keep-original "first_stage_model.*"
120
+ ```
121
+
122
+ **Turn a pickled checkpoint into safetensors without changing precision:**
123
+
124
+ ```bash
125
+ precision-convert model.ckpt -st
126
+ ```
127
+
128
+ **Check what's inside a file before converting:**
129
+
130
+ ```console
131
+ $ precision-convert model.ckpt --info
132
+ model.ckpt: 523 tensors, 27,274,319 values, 109.10 MB
133
+ float32 521 tensors 109.10 MB
134
+ int64 1 tensor 616 B
135
+ int32 1 tensor 4 B
136
+ EMA: 242 model_ema.* tensors (-t ema-only or -t no-ema removes them)
137
+ ```
138
+
139
+ ## Which precision should I pick?
140
+
141
+ | Precision | Bytes per value | Largest value | Good for |
142
+ |---|---|---|---|
143
+ | `fp32` | 4 | 3.4 × 10³⁸ | Training, reference copies, upcasting half-precision files |
144
+ | `bf16` | 2 | 3.4 × 10³⁸ | The safe 2-byte choice: same range as fp32, so nothing overflows |
145
+ | `fp16` | 2 | 65,504 | GPUs without fast bf16, and tools that expect fp16 |
146
+ | `fp8_e4m3fn` | 1 | 448 | Smallest files, for inference in tools that load fp8 weights; some quality loss |
147
+ | `fp8_e5m2` | 1 | 57,344 | fp8 with more range and less precision than `e4m3fn` |
148
+
149
+ The fp8 options cast each value directly, without per-tensor scaling factors.
150
+
151
+ ## Python API
152
+
153
+ ```python
154
+ from precision_converter import convert_state_dict, load_checkpoint, save_checkpoint
155
+
156
+ state_dict, metadata = load_checkpoint("model.safetensors")
157
+ converted, report = convert_state_dict(state_dict, precision="bf16", conversion="no-ema")
158
+ save_checkpoint(converted, "model-bf16.safetensors", metadata=metadata)
159
+
160
+ print(f"{report.converted} tensors converted, {len(report.clamped)} had values clamped")
161
+ ```
162
+
163
+ `convert_state_dict` takes any mapping of tensors, so `model.state_dict()` works too.
164
+
165
+ ## Upgrading from 1.x
166
+
167
+ `python convert.py ...` still works, but 2.0 fixes bugs that changed what earlier versions wrote:
168
+
169
+ - `-t ema-only` used to keep the regular weights and store the EMA weights under keys no model loads. It now
170
+ outputs the actual EMA weights.
171
+ - Integer tensors are no longer cast to floating point, and overflowing values are clamped instead of
172
+ becoming `inf`.
173
+ - The default precision is now `keep`. `-p fp32` used to change nothing; it now converts fp16 and bf16 tensors
174
+ to fp32. With the default precision the output name no longer ends in `-fp32`.
175
+ - Checkpoints are loaded with `weights_only=True`. Use `--unsafe-load` for trusted files that need more.
176
+ - Existing output files are no longer overwritten unless you pass `--overwrite`.
177
+
178
+ See [CHANGELOG.md](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/CHANGELOG.md) for details.
179
+
180
+ ## Contributing
181
+
182
+ Issues and pull requests are welcome. To run the tests:
183
+
184
+ ```bash
185
+ pip install -e ".[test]"
186
+ pytest
187
+ ```
188
+
189
+ ## License
190
+
191
+ MIT. See [LICENSE](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE).
@@ -0,0 +1,165 @@
1
+ # PyTorch Precision Converter
2
+
3
+ [![PyPI](https://img.shields.io/pypi/v/pytorch-precision-converter)](https://pypi.org/project/pytorch-precision-converter/)
4
+ [![Tests](https://github.com/angelolamonaca/PyTorch-Precision-Converter/actions/workflows/ci.yml/badge.svg)](https://github.com/angelolamonaca/PyTorch-Precision-Converter/actions/workflows/ci.yml)
5
+ ![Python](https://img.shields.io/badge/python-3.9%2B-blue)
6
+ ![PyTorch](https://img.shields.io/badge/PyTorch-2.1%2B-orange)
7
+ [![License: MIT](https://img.shields.io/badge/license-MIT-green)](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE)
8
+
9
+ Make PyTorch checkpoints and `safetensors` files **half the size (a quarter with fp8)** by converting them to
10
+ `fp16`, `bf16` or `fp8`, and strip the EMA weights you don't need for inference. One command, no model code
11
+ needed, runs on the CPU.
12
+
13
+ ```console
14
+ $ precision-convert model.ckpt -p fp16 -t ema-only -st
15
+ Converted 280 tensors to fp16, left 1 unchanged, dropped 242 EMA entries
16
+ Used the EMA weights for 240 tensors
17
+ Saved model-ema-only-fp16.safetensors (29.93 MB, 27% of 109.23 MB)
18
+ ```
19
+
20
+ ## Why not just call `.half()`?
21
+
22
+ A one-line cast quietly breaks things. This tool handles the details:
23
+
24
+ - **Only floating-point tensors are cast.** Integer and bool tensors like `position_ids` or step counters keep
25
+ their dtype instead of turning into floats.
26
+ - **No silent `inf` or `NaN`.** fp16 tops out at 65,504 and fp8 at 448. Values beyond the range are clamped
27
+ to it and you get a warning, rather than infinities appearing in your weights.
28
+ - **Real EMA extraction.** `-t ema-only` puts the EMA weights where the model loads them from, for Stable
29
+ Diffusion and other latent-diffusion checkpoints (`model_ema.*`), then drops the duplicates.
30
+ - **safetensors metadata survives.** LoRA training info (`ss_*` keys) and other metadata is copied to the output.
31
+ - **Safe by default.** PyTorch checkpoints load with `weights_only=True`, so a malicious `.ckpt` can't run code.
32
+ Files that need full unpickling require an explicit `--unsafe-load`.
33
+ - **Your files are safe too.** Output is written to a temporary file and moved into place when complete, the
34
+ input is never replaced, and existing files are only overwritten with `--overwrite`.
35
+ - **Mixed precision when you need it.** `--keep-original "first_stage_model.*"` keeps matching tensors (here the
36
+ VAE) in their original dtype.
37
+
38
+ ## Install
39
+
40
+ ```bash
41
+ pip install pytorch-precision-converter
42
+ ```
43
+
44
+ For the latest unreleased changes, install from GitHub instead:
45
+ `pip install git+https://github.com/angelolamonaca/PyTorch-Precision-Converter`.
46
+
47
+ Either way you get the `precision-convert` command (`python -m precision_converter` works too). It needs Python 3.9+,
48
+ PyTorch 2.1+ and `safetensors`.
49
+
50
+ ## Usage
51
+
52
+ ```bash
53
+ precision-convert INPUT [-p PRECISION] [-t TYPE] [-st] [-o OUTPUT] [options]
54
+ ```
55
+
56
+ | Option | What it does |
57
+ |---|---|
58
+ | `INPUT` | `.safetensors` file or PyTorch checkpoint (`.ckpt`, `.pt`, `.pth`, `.bin`). `-f INPUT` works too. |
59
+ | `-p`, `--precision` | `fp32`, `fp16`, `bf16`, `fp8_e4m3fn`, `fp8_e5m2`, or `keep` (default: change no dtypes). |
60
+ | `-t`, `--type` | `full` keeps everything (default). `ema-only` uses the EMA weights. `no-ema` drops the EMA weights. |
61
+ | `-st`, `--safe-tensors` | Save as `.safetensors` instead of a PyTorch `.ckpt`. |
62
+ | `-o`, `--output` | Output path; its extension picks the format. Default: `<input>-<type>-<precision>` next to the input. |
63
+ | `--keep-original PATTERN` | Leave tensors whose key matches the glob in their original dtype. Repeatable. |
64
+ | `--info` | Show the tensors, dtypes, sizes, EMA and metadata in the file, then exit. |
65
+ | `--unsafe-load` | Fully unpickle PyTorch checkpoints. Only for files you trust. |
66
+ | `--overwrite` | Replace the output file if it exists. |
67
+ | `-v`, `--verbose` | List each tensor that had values clamped. |
68
+
69
+ ## Recipes
70
+
71
+ **Shrink a Stable Diffusion training checkpoint for inference** (EMA weights, fp16, safetensors):
72
+
73
+ ```bash
74
+ precision-convert model.ckpt -p fp16 -t ema-only -st
75
+ ```
76
+
77
+ **Convert a LoRA to fp16 and keep its training metadata:**
78
+
79
+ ```bash
80
+ precision-convert my-lora.safetensors -p fp16 -o my-lora-fp16.safetensors
81
+ ```
82
+
83
+ **Store a large model in fp8** for UIs that load fp8 weights, such as ComfyUI, keeping normalization layers
84
+ precise:
85
+
86
+ ```bash
87
+ precision-convert model.safetensors -p fp8_e4m3fn -st --keep-original "*norm*"
88
+ ```
89
+
90
+ **Keep the VAE in fp32** (fp16 VAEs can produce NaNs and black images in some models):
91
+
92
+ ```bash
93
+ precision-convert model.safetensors -p fp16 -st --keep-original "first_stage_model.*"
94
+ ```
95
+
96
+ **Turn a pickled checkpoint into safetensors without changing precision:**
97
+
98
+ ```bash
99
+ precision-convert model.ckpt -st
100
+ ```
101
+
102
+ **Check what's inside a file before converting:**
103
+
104
+ ```console
105
+ $ precision-convert model.ckpt --info
106
+ model.ckpt: 523 tensors, 27,274,319 values, 109.10 MB
107
+ float32 521 tensors 109.10 MB
108
+ int64 1 tensor 616 B
109
+ int32 1 tensor 4 B
110
+ EMA: 242 model_ema.* tensors (-t ema-only or -t no-ema removes them)
111
+ ```
112
+
113
+ ## Which precision should I pick?
114
+
115
+ | Precision | Bytes per value | Largest value | Good for |
116
+ |---|---|---|---|
117
+ | `fp32` | 4 | 3.4 × 10³⁸ | Training, reference copies, upcasting half-precision files |
118
+ | `bf16` | 2 | 3.4 × 10³⁸ | The safe 2-byte choice: same range as fp32, so nothing overflows |
119
+ | `fp16` | 2 | 65,504 | GPUs without fast bf16, and tools that expect fp16 |
120
+ | `fp8_e4m3fn` | 1 | 448 | Smallest files, for inference in tools that load fp8 weights; some quality loss |
121
+ | `fp8_e5m2` | 1 | 57,344 | fp8 with more range and less precision than `e4m3fn` |
122
+
123
+ The fp8 options cast each value directly, without per-tensor scaling factors.
124
+
125
+ ## Python API
126
+
127
+ ```python
128
+ from precision_converter import convert_state_dict, load_checkpoint, save_checkpoint
129
+
130
+ state_dict, metadata = load_checkpoint("model.safetensors")
131
+ converted, report = convert_state_dict(state_dict, precision="bf16", conversion="no-ema")
132
+ save_checkpoint(converted, "model-bf16.safetensors", metadata=metadata)
133
+
134
+ print(f"{report.converted} tensors converted, {len(report.clamped)} had values clamped")
135
+ ```
136
+
137
+ `convert_state_dict` takes any mapping of tensors, so `model.state_dict()` works too.
138
+
139
+ ## Upgrading from 1.x
140
+
141
+ `python convert.py ...` still works, but 2.0 fixes bugs that changed what earlier versions wrote:
142
+
143
+ - `-t ema-only` used to keep the regular weights and store the EMA weights under keys no model loads. It now
144
+ outputs the actual EMA weights.
145
+ - Integer tensors are no longer cast to floating point, and overflowing values are clamped instead of
146
+ becoming `inf`.
147
+ - The default precision is now `keep`. `-p fp32` used to change nothing; it now converts fp16 and bf16 tensors
148
+ to fp32. With the default precision the output name no longer ends in `-fp32`.
149
+ - Checkpoints are loaded with `weights_only=True`. Use `--unsafe-load` for trusted files that need more.
150
+ - Existing output files are no longer overwritten unless you pass `--overwrite`.
151
+
152
+ See [CHANGELOG.md](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/CHANGELOG.md) for details.
153
+
154
+ ## Contributing
155
+
156
+ Issues and pull requests are welcome. To run the tests:
157
+
158
+ ```bash
159
+ pip install -e ".[test]"
160
+ pytest
161
+ ```
162
+
163
+ ## License
164
+
165
+ MIT. See [LICENSE](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE).
@@ -0,0 +1,27 @@
1
+ """Convert the tensor precision of PyTorch checkpoints and safetensors files."""
2
+
3
+ __version__ = "2.0.0"
4
+
5
+ from .core import ( # noqa: E402
6
+ CONVERSION_TYPES,
7
+ PRECISIONS,
8
+ CheckpointInfo,
9
+ ConversionReport,
10
+ UnsafeCheckpointError,
11
+ convert_state_dict,
12
+ describe,
13
+ load_checkpoint,
14
+ save_checkpoint,
15
+ )
16
+
17
+ __all__ = [
18
+ "CONVERSION_TYPES",
19
+ "PRECISIONS",
20
+ "CheckpointInfo",
21
+ "ConversionReport",
22
+ "UnsafeCheckpointError",
23
+ "convert_state_dict",
24
+ "describe",
25
+ "load_checkpoint",
26
+ "save_checkpoint",
27
+ ]
@@ -0,0 +1,5 @@
1
+ import sys
2
+
3
+ from .cli import main
4
+
5
+ sys.exit(main())
@@ -0,0 +1,181 @@
1
+ """Command-line interface: ``precision-convert``."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import sys
7
+ from pathlib import Path
8
+
9
+ import torch
10
+
11
+ from . import __version__
12
+ from .core import (
13
+ CONVERSION_TYPES,
14
+ PRECISIONS,
15
+ UnsafeCheckpointError,
16
+ convert_state_dict,
17
+ describe,
18
+ load_checkpoint,
19
+ save_checkpoint,
20
+ )
21
+
22
+ EXAMPLES = """\
23
+ examples:
24
+ precision-convert model.safetensors -p fp16 -st
25
+ precision-convert sd-model.ckpt -p fp16 -t ema-only -st
26
+ precision-convert lora.safetensors -p bf16 -o lora-bf16.safetensors
27
+ precision-convert flux.safetensors -p fp8_e4m3fn -st --keep-original "*norm*"
28
+ precision-convert model.ckpt --info
29
+ """
30
+
31
+
32
+ def build_parser() -> argparse.ArgumentParser:
33
+ parser = argparse.ArgumentParser(
34
+ prog="precision-convert",
35
+ description="Convert the tensor precision of PyTorch checkpoints and safetensors files.",
36
+ epilog=EXAMPLES,
37
+ formatter_class=argparse.RawDescriptionHelpFormatter,
38
+ )
39
+ parser.add_argument(
40
+ "input", nargs="?", help="Model file: .safetensors or a PyTorch checkpoint (.ckpt, .pt, .pth, .bin)."
41
+ )
42
+ parser.add_argument("-f", "--file", help="Same as INPUT, kept for compatibility with 1.x.")
43
+ parser.add_argument(
44
+ "-p",
45
+ "--precision",
46
+ default="keep",
47
+ choices=PRECISIONS,
48
+ help="Target precision for floating-point tensors (default: keep, which changes no dtype).",
49
+ )
50
+ parser.add_argument(
51
+ "-t",
52
+ "--type",
53
+ default="full",
54
+ choices=CONVERSION_TYPES,
55
+ help="full: keep everything (default). ema-only: use the EMA weights and drop the EMA copies. "
56
+ "no-ema: drop the EMA copies.",
57
+ )
58
+ parser.add_argument(
59
+ "-st", "--safe-tensors", action="store_true", help="Save in safetensors format (default: PyTorch .ckpt)."
60
+ )
61
+ parser.add_argument(
62
+ "-o",
63
+ "--output",
64
+ help="Output path. Its extension picks the format (.safetensors or anything else for PyTorch). "
65
+ "Default: <input>-<type>-<precision> next to the input.",
66
+ )
67
+ parser.add_argument(
68
+ "--keep-original",
69
+ metavar="PATTERN",
70
+ action="append",
71
+ default=[],
72
+ help='Leave tensors whose key matches this glob in their original dtype, e.g. "first_stage_model.*". '
73
+ "Repeatable.",
74
+ )
75
+ parser.add_argument("--overwrite", action="store_true", help="Replace the output file if it already exists.")
76
+ parser.add_argument(
77
+ "--unsafe-load",
78
+ action="store_true",
79
+ help="Load PyTorch checkpoints with full unpickling (weights_only=False). Unpickling can run code "
80
+ "hidden in the file: only use this for files you trust.",
81
+ )
82
+ parser.add_argument("--info", action="store_true", help="Describe the input file and exit without converting.")
83
+ parser.add_argument("-v", "--verbose", action="store_true", help="List the tensors whose values were clamped.")
84
+ parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
85
+ return parser
86
+
87
+
88
+ def format_bytes(n: float) -> str:
89
+ for unit in ("B", "KB", "MB", "GB"):
90
+ if n < 1000 or unit == "GB":
91
+ return f"{n:.0f} {unit}" if unit == "B" else f"{n:.2f} {unit}"
92
+ n /= 1000
93
+ raise AssertionError("unreachable")
94
+
95
+
96
+ def default_output(input_path: Path, conversion: str, precision: str, safetensors: bool) -> Path:
97
+ parts = [input_path.stem, conversion] + ([precision] if precision != "keep" else [])
98
+ return input_path.with_name("-".join(parts) + (".safetensors" if safetensors else ".ckpt"))
99
+
100
+
101
+ def print_info(path: Path, state_dict: dict, metadata: dict | None) -> None:
102
+ info = describe(state_dict)
103
+ print(f"{path}: {info.tensors:,} tensors, {info.parameters:,} values, {format_bytes(info.bytes)}")
104
+ for dtype, (count, nbytes) in sorted(info.dtypes.items(), key=lambda item: -item[1][1]):
105
+ print(f" {dtype:<16} {count:>7,} {'tensor ' if count == 1 else 'tensors'} {format_bytes(nbytes):>10}")
106
+ if info.ema_tensors:
107
+ print(f" EMA: {info.ema_tensors:,} model_ema.* tensors (-t ema-only or -t no-ema removes them)")
108
+ if info.other_keys:
109
+ print(f" Non-tensor keys: {', '.join(info.other_keys)}")
110
+ if metadata:
111
+ print(f" Metadata: {len(metadata)} entries ({', '.join(sorted(metadata)[:8])}{', ...' if len(metadata) > 8 else ''})")
112
+
113
+
114
+ def main(argv: list[str] | None = None) -> int:
115
+ parser = build_parser()
116
+ args = parser.parse_args(argv)
117
+
118
+ if args.input and args.file:
119
+ parser.error("give the input file either as INPUT or with -f, not both")
120
+ input_path = Path(args.input or args.file or "model.ckpt")
121
+ if not input_path.is_file():
122
+ parser.error(f"input file not found: {input_path}")
123
+
124
+ if args.output:
125
+ output_path = Path(args.output)
126
+ if args.safe_tensors and output_path.suffix != ".safetensors":
127
+ parser.error("--safe-tensors needs an --output path ending in .safetensors")
128
+ else:
129
+ output_path = default_output(input_path, args.type, args.precision, args.safe_tensors)
130
+
131
+ try:
132
+ state_dict, metadata = load_checkpoint(input_path, unsafe=args.unsafe_load)
133
+ except (UnsafeCheckpointError, ValueError, OSError, RuntimeError) as e:
134
+ print(f"error: {e}", file=sys.stderr)
135
+ return 1
136
+
137
+ if args.info:
138
+ print_info(input_path, state_dict, metadata)
139
+ return 0
140
+
141
+ if output_path.resolve() == input_path.resolve():
142
+ parser.error("the output would replace the input file; choose another --output")
143
+ if output_path.exists() and not args.overwrite:
144
+ parser.error(f"{output_path} already exists; pass --overwrite to replace it")
145
+
146
+ converted, report = convert_state_dict(state_dict, args.precision, args.type, args.keep_original)
147
+
148
+ if args.precision == "keep":
149
+ summary = f"Kept every dtype as it was ({report.unchanged:,} entries)"
150
+ else:
151
+ summary = f"Converted {report.converted:,} tensors to {args.precision}, left {report.unchanged:,} unchanged"
152
+ if report.dropped:
153
+ summary += f", dropped {report.dropped:,} EMA entries"
154
+ print(summary)
155
+ if args.type == "ema-only":
156
+ if report.ema_swapped:
157
+ print(f"Used the EMA weights for {report.ema_swapped:,} tensors")
158
+ else:
159
+ print("warning: no EMA weights found; the output keeps the regular weights", file=sys.stderr)
160
+ if report.clamped:
161
+ limit = torch.finfo(PRECISIONS[args.precision]).max
162
+ print(
163
+ f"warning: {len(report.clamped):,} tensors had values outside the {args.precision} range "
164
+ f"(±{limit:g}); they were clamped instead of becoming inf/NaN",
165
+ file=sys.stderr,
166
+ )
167
+ if args.verbose:
168
+ for key, count in report.clamped.items():
169
+ print(f" {key}: {count:,} values", file=sys.stderr)
170
+
171
+ skipped = save_checkpoint(converted, output_path, metadata=metadata)
172
+ if skipped:
173
+ print(f"warning: safetensors only stores tensors, left out: {', '.join(skipped)}", file=sys.stderr)
174
+
175
+ before, after = input_path.stat().st_size, output_path.stat().st_size
176
+ print(f"Saved {output_path} ({format_bytes(after)}, {after / before:.0%} of {format_bytes(before)})")
177
+ return 0
178
+
179
+
180
+ if __name__ == "__main__":
181
+ sys.exit(main())