Speed up the diffusers example tests (#2724)

### What does this PR do?

Type of change: Tests (CI speed-up)

The diffusers example lane spends ~29 min in pytest, almost all of it in
~45 example-script subprocesses on tiny models. This PR removes
redundant and avoidable work without dropping coverage:

- **Export HF checkpoints from runs that already happen.** The SDXL
int8/fp8 and Wan 2.2 14B transformer int8/fp8 configs in
`test_export_diffusers_hf_ckpt.py` used exactly the same arguments as
runs in `test_diffusers.py`. Those quantize runs now also pass
`--hf-ckpt-dir` (`hf_export=True`) and check the export, and the four
duplicate export tests are removed (30 → 26 tests in these files).
- **Calibrate from a local prompts file.** `quantize.py` loaded the Hub
prompt datasets on every calibration run. For Wan 2.2 that is all of
`nkp37/OpenVid-1M`, which cost ~30s per run. Tests now pass
`--prompts-file tests/examples/diffusers/calib_prompts.txt`, which also
removes a network dependency.
- **Tiny PixArt for `test_pixart_cachify`.** It downloaded the full-size
`PixArt-XL-2-1024-MS`. It now uses a tiny local PixArt-Alpha pipeline,
with the same tiny config as diffusers' own PixArt tests and built like
the existing tiny Wan 2.2 fixture. `cachify` still wraps both
transformer blocks. `PIXART_PATH` is removed because nothing else uses
it.
- **Cheaper `torch.compile` tests.** `diffusion_trt.py` compiles with
`mode="max-autotune"`. The compile tests now set
`TORCHINDUCTOR_MAX_AUTOTUNE_{GEMM,CONV}_BACKENDS=ATEN`, so Inductor
skips benchmarking Triton GEMM/conv templates. The `--torch-compile`
path still runs.

### Usage

N/A (tests only).

### Testing

I ran these on an RTX PRO 6000 (same GPU as the CI runner) in
`nvcr.io/nvidia/pytorch:26.07-py3`, set up like the CI lane: editable
`.[onnx,hf,dev-test]`, `diffusers==0.39.0`, apex uninstalled. Coverage
was on, the same way CI collects it (`COVERAGE_PROCESS_START`, `--cov`).

- Branch: `pytest tests/examples/diffusers
--ignore=test_diffusers_onnx_trt.py --ignore=fastgen --cov` → **40
passed in 22m37s**. fastgen needs `nemo_automodel`, which I didn't
install; this PR doesn't touch it.
- `main`, same machine: `test_diffusers.py`,
`test_export_diffusers_hf_ckpt.py`, `test_cache_diffusion.py -k "not
pixart"` → **29 passed in 28m21s**. This set is smaller than the branch
run: no sparsity, PixArt or deprecation tests (~1m50s on the branch).

Per-test call times (s), with CI from run 37945140337:

| Test | CI (main) | local main | local branch |
|---|---:|---:|---:|
| `test_wan22_quantization[*]` (7 tests, each) | 66–86 | 79–98 | 49–52 |
| `test_wan22_hf_ckpt_export` (2 tests) | 44 + 46 | 54 + 54 | removed |
| `test_diffusers_hf_ckpt_export[sdxl int8 / fp8]` | 43 + 37 | 32 + 31 |
removed |
| `test_pixart_cachify` | 86 | not run | 1.8 |
| `test_diffusion_trt_torch[sdxl_1.0_torch_compile]` | 132 | 156 | 119 |
| `test_diffusion_trt_torch[flux_schnell_torch_compile]` | 60 | 71 | 57
|
| `test_diffusion_trt_torch[sd3_medium_torch_compile]` | 51 | 58 | 49 |

On the same test set this saves about 7.5 min locally, plus about 84s
for PixArt. Scaled to CI timings, I expect the lane's pytest step to
drop from ~29 min to ~21 min.

The ATen restriction cuts the SDXL compile run from 108s to 63s without
coverage (run on its own). Coverage tracing of Inductor's Python
accounts for most of what remains.

`pre-commit` passes on the changed files.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: N/A (tests only)
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A (the tiny
PixArt config values mirror diffusers' own tests, as the existing tiny
Wan 2.2 helper does)
- Did you write any new necessary tests?: N/A
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A (tests only)
- Did you get Claude approval on this PR?: ❌

### Additional Information

Possible follow-ups not in this PR:
- Run fewer restore subprocesses: one per format instead of all 13,
about 3 min, but it gives up some coverage.
- Run the example scripts in-process like `megatron_example_runner.py`:
about 20s of startup saved per subprocess, but a larger change.

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Tests**
* Expanded diffusion-model test coverage for PixArt-Alpha, including
cache behavior and checkpoint exports.
* Added validation that exported checkpoints include configuration and
model-weight files.
* Updated quantization tests to use local calibration prompts and cover
checkpoint export for selected SDXL and Wan 2.2 runs.
  * Adjusted compiled test runs to use compatible autotuning settings.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

tests/_test_utils/examples/models.py

@@ -58,8 +58,3 @@ FLUX_SCHNELL_PATH = _select_path(
     remote_id="hf-internal-testing/tiny-flux-pipe",
     local_id="FLUX.1-schnell",
 )
-
-PIXART_PATH = _select_path(
-    remote_id="PixArt-alpha/PixArt-XL-2-1024-MS",
-    local_id="PixArt-XL-2-1024-MS",
-)

tests/_test_utils/torch/diffusers_models.py

@@ -263,6 +263,56 @@ def create_tiny_wan22_pipeline_dir(tmp_path: Path) -> Path:
     return save_dir
 
 
+def create_tiny_pixart_pipeline_dir(tmp_path: Path) -> Path:
+    """Create and save a tiny PixArt-Alpha pipeline to a directory.
+
+    Uses the same tiny config as diffusers' own PixArt-Alpha tests (2 transformer blocks,
+    hf-internal-testing/tiny-random-t5 text encoder). Its ``sample_size`` has no resolution
+    bin, so call the pipeline with ``use_resolution_binning=False``.
+
+    The saved directory can be loaded with ``PixArtAlphaPipeline.from_pretrained(path)``.
+    """
+    from diffusers import (
+        AutoencoderKL,
+        DDIMScheduler,
+        PixArtAlphaPipeline,
+        PixArtTransformer2DModel,
+    )
+    from transformers import AutoTokenizer, T5EncoderModel
+
+    torch.manual_seed(0)
+    transformer = PixArtTransformer2DModel(
+        sample_size=8,
+        num_layers=2,
+        patch_size=2,
+        attention_head_dim=8,
+        num_attention_heads=3,
+        caption_channels=32,
+        in_channels=4,
+        cross_attention_dim=24,
+        out_channels=8,
+        attention_bias=True,
+        activation_fn="gelu-approximate",
+        num_embeds_ada_norm=1000,
+        norm_type="ada_norm_single",
+        norm_elementwise_affine=False,
+        norm_eps=1e-6,
+    )
+
+    torch.manual_seed(0)
+    pipe = PixArtAlphaPipeline(
+        transformer=transformer,
+        vae=AutoencoderKL(),
+        scheduler=DDIMScheduler(),
+        text_encoder=T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"),
+        tokenizer=AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"),
+    )
+
+    save_dir = tmp_path / "tiny_pixart"
+    pipe.save_pretrained(save_dir)
+    return save_dir
+
+
 def get_tiny_qwen_image_transformer(**config_kwargs):
     """Create a tiny QwenImageTransformer2DModel for testing.
 

tests/examples/diffusers/calib_prompts.txt

@@ -0,0 +1,4 @@
+a small cactus with a happy face in the Sahara desert
+a photo of an astronaut riding a horse on the moon
+a lighthouse on a rocky cliff at sunset, oil painting
+a red fox walking through a snowy forest, cinematic lighting

tests/examples/diffusers/conftest.py

@@ -15,7 +15,10 @@
 
 import pytest
 from _test_utils.fs_utils import assert_unmodified_tree
-from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir
+from _test_utils.torch.diffusers_models import (
+    create_tiny_pixart_pipeline_dir,
+    create_tiny_qwen_image_pipeline_dir,
+)
 
 
 @pytest.fixture(scope="session")
@@ -23,3 +26,10 @@ def tiny_qwen_image_path(tmp_path_factory):
     tmp_path = tmp_path_factory.mktemp("qwen_image")
     with assert_unmodified_tree(create_tiny_qwen_image_pipeline_dir(tmp_path)) as path:
         yield str(path)
+
+
+@pytest.fixture(scope="session")
+def tiny_pixart_path(tmp_path_factory):
+    tmp_path = tmp_path_factory.mktemp("pixart")
+    with assert_unmodified_tree(create_tiny_pixart_pipeline_dir(tmp_path)) as path:
+        yield str(path)

tests/examples/diffusers/test_cache_diffusion.py

@@ -18,7 +18,7 @@ import sys
 
 import pytest
 import torch
-from _test_utils.examples.models import PIXART_PATH, SDXL_PATH
+from _test_utils.examples.models import SDXL_PATH
 from _test_utils.examples.run_command import MODELOPT_ROOT
 from diffusers import DiffusionPipeline, PixArtAlphaPipeline
 
@@ -45,18 +45,25 @@ def test_sdxl_cachify():
     torch.cuda.empty_cache()
 
 
-def test_pixart_cachify():
+def test_pixart_cachify(tiny_pixart_path):
     # Fail test if apex is installed
     if "apex" in subprocess.check_output(["pip", "list"]).decode("utf-8"):
         pytest.xfail("Apex is installed, test is expected to fail")
 
-    pipe = PixArtAlphaPipeline.from_pretrained(PIXART_PATH, torch_dtype=torch.float16).to("cuda")
+    pipe = PixArtAlphaPipeline.from_pretrained(tiny_pixart_path, torch_dtype=torch.float16).to(
+        "cuda"
+    )
     cachify.prepare(pipe, PIXART_DEFAULT_CONFIG)
 
     prompt = "a small cactus with a happy face in the Sahara desert"
     generator = torch.Generator(device="cuda").manual_seed(2946901)
     # 8 steps still exercises the step-modulo cache pattern; this is a runs-without-error smoke test.
-    pipe(prompt=prompt, generator=generator, num_inference_steps=8).images[0]
+    pipe(
+        prompt=prompt,
+        generator=generator,
+        num_inference_steps=8,
+        use_resolution_binning=False,  # the tiny transformer's sample_size has no resolution bin
+    ).images[0]
     # Clear cuda memory as pytest doesnt clear it between tests
     del pipe
     torch.cuda.empty_cache()

tests/examples/diffusers/test_diffusers.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from pathlib import Path
 from typing import NamedTuple
 
@@ -21,6 +22,20 @@ from _test_utils.examples.models import FLUX_SCHNELL_PATH, SD3_PATH, SDXL_PATH
 from _test_utils.examples.run_command import run_example_command
 from _test_utils.torch.misc import minimum_sm
 
+# Calibrate from local prompts: the default Hub datasets (OpenVid-1M for Wan) are slow to load
+# and add a network dependency without adding coverage.
+CALIB_PROMPTS_FILE = Path(__file__).parent / "calib_prompts.txt"
+
+
+def assert_hf_ckpt_exported(hf_ckpt_dir: Path) -> None:
+    assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}"
+
+    config_files = list(hf_ckpt_dir.rglob("config.json"))
+    assert len(config_files) > 0, f"No config.json found in {hf_ckpt_dir}"
+
+    weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin"))
+    assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}"
+
 
 class DiffuserModel(NamedTuple):
     dtype: str
@@ -29,6 +44,9 @@ class DiffuserModel(NamedTuple):
     format_type: str
     quant_algo: str
     collect_method: str
+    # Also export an HF checkpoint from the quantize run, rather than calibrating again in a
+    # separate export test.
+    hf_export: bool = False
 
     def _run_cmd(self, script: str, *args: str) -> str:
         cmd_args = [
@@ -60,6 +78,8 @@ class DiffuserModel(NamedTuple):
             self.collect_method,
             "--quant-algo",
             self.quant_algo,
+            "--prompts-file",
+            str(CALIB_PROMPTS_FILE),
         ]
 
     def quantize(self, tmp_path: Path, *export_args: str) -> None:
@@ -94,6 +114,8 @@ class Wan22Model(NamedTuple):
     format_type: str
     quant_algo: str
     collect_method: str
+    # See DiffuserModel.hf_export.
+    hf_export: bool = False
 
     def _ckpt_path(self, tmp_path: Path) -> str:
         stem = self.model.replace("wan2.2-t2v-", "")
@@ -124,6 +146,8 @@ class Wan22Model(NamedTuple):
             "1",
             "--n-steps",
             "2",
+            "--prompts-file",
+            str(CALIB_PROMPTS_FILE),
             # Tiny video dims — override MODEL_DEFAULTS for fast CI.
             "--extra-param",
             "height=16",
@@ -136,12 +160,13 @@ class Wan22Model(NamedTuple):
             cmd_args.extend(["--backbone", self.backbone])
         return cmd_args
 
-    def quantize(self, tiny_wan22_path: str, tmp_path: Path) -> None:
+    def quantize(self, tiny_wan22_path: str, tmp_path: Path, *export_args: str) -> None:
         run_example_command(
             [
                 *self._common_args(tiny_wan22_path),
                 "--quantized-torch-ckpt-save-path",
                 self._ckpt_path(tmp_path),
+                *export_args,
             ],
             "diffusers/quantization",
         )
@@ -196,6 +221,7 @@ DIFFUSER_MODELS = [
             format_type="fp8",
             quant_algo="max",
             collect_method="default",
+            hf_export=True,
         ),
         marks=minimum_sm(89),
         id="sdxl_1.0_fp16_fp8_max_3.0_default",
@@ -220,6 +246,7 @@ DIFFUSER_MODELS = [
             format_type="int8",
             quant_algo="smoothquant",
             collect_method="min-mean",
+            hf_export=True,
         ),
         id="sdxl_1.0_fp16_int8_smoothquant_3.0_min_mean",
     ),
@@ -232,11 +259,11 @@ DIFFUSER_MODELS = [
 # for each.
 WAN22_MODELS = [
     pytest.param(
-        Wan22Model("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"),
+        Wan22Model("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean", hf_export=True),
         id="wan22_14b_transformer_int8_smoothquant",
     ),
     pytest.param(
-        Wan22Model("wan2.2-t2v-14b", None, "fp8", "max", "default"),
+        Wan22Model("wan2.2-t2v-14b", None, "fp8", "max", "default", hf_export=True),
         marks=minimum_sm(89),
         id="wan22_14b_transformer_fp8_max",
     ),
@@ -270,13 +297,23 @@ WAN22_MODELS = [
 
 @pytest.mark.parametrize("model", DIFFUSER_MODELS)
 def test_diffusers_quantization(model: DiffuserModel, tmp_path: Path) -> None:
-    model.quantize(tmp_path)
+    hf_ckpt_dir = tmp_path / "hf_ckpt"
+    model.quantize(tmp_path, *(["--hf-ckpt-dir", str(hf_ckpt_dir)] if model.hf_export else []))
+    if model.hf_export:
+        assert_hf_ckpt_exported(hf_ckpt_dir)
     model.restore(tmp_path)
 
 
 @pytest.mark.parametrize("wan_model", WAN22_MODELS)
 def test_wan22_quantization(wan_model: Wan22Model, tiny_wan22_path: str, tmp_path: Path) -> None:
-    wan_model.quantize(tiny_wan22_path, tmp_path)
+    hf_ckpt_dir = tmp_path / "hf_ckpt"
+    wan_model.quantize(
+        tiny_wan22_path,
+        tmp_path,
+        *(["--hf-ckpt-dir", str(hf_ckpt_dir)] if wan_model.hf_export else []),
+    )
+    if wan_model.hf_export:
+        assert_hf_ckpt_exported(hf_ckpt_dir)
     wan_model.restore(tiny_wan22_path, tmp_path)
 
 
@@ -315,6 +352,14 @@ def test_diffusion_trt_torch(
         "--num-inference-steps",
         "2",
     ]
+    env = None
     if torch_compile:
         cmd_args.append("--torch-compile")
-    run_example_command(cmd_args, "diffusers/quantization")
+        # The script compiles with mode="max-autotune"; benchmarking Triton GEMM/conv templates is
+        # a large share of the run on these tiny models and covers nothing in the example.
+        env = {
+            **os.environ,
+            "TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS": "ATEN",
+            "TORCHINDUCTOR_MAX_AUTOTUNE_CONV_BACKENDS": "ATEN",
+        }
+    run_example_command(cmd_args, "diffusers/quantization", env=env)

tests/examples/diffusers/test_export_diffusers_hf_ckpt.py

@@ -19,10 +19,11 @@ from pathlib import Path
 from typing import NamedTuple
 
 import pytest
-from _test_utils.examples.models import FLUX_SCHNELL_PATH, SDXL_PATH
+from _test_utils.examples.models import FLUX_SCHNELL_PATH
 from _test_utils.examples.run_command import run_example_command
 from _test_utils.torch.misc import minimum_sm
 from safetensors import safe_open
+from test_diffusers import CALIB_PROMPTS_FILE, assert_hf_ckpt_exported
 
 
 class DiffuserHfExportModel(NamedTuple):
@@ -63,6 +64,8 @@ class DiffuserHfExportModel(NamedTuple):
             self.model_dtype,
             "--trt-high-precision-dtype",
             self.dtype,
+            "--prompts-file",
+            str(CALIB_PROMPTS_FILE),
             "--hf-ckpt-dir",
             str(hf_ckpt_dir),
         ]
@@ -70,17 +73,11 @@ class DiffuserHfExportModel(NamedTuple):
         return hf_ckpt_dir
 
 
+# The SDXL int8/fp8 and Wan 2.2 transformer int8/fp8 exports run in test_diffusers.py, from the
+# quantize run those tests already do (``hf_export=True``).
 @pytest.mark.parametrize(
     "model",
     [
-        DiffuserHfExportModel(
-            name="sdxl-1.0",
-            path=SDXL_PATH,
-            dtype="Half",
-            format_type="int8",
-            quant_algo="smoothquant",
-            collect_method="min-mean",
-        ),
         DiffuserHfExportModel(
             name="flux-schnell",
             path=FLUX_SCHNELL_PATH,
@@ -90,17 +87,6 @@ class DiffuserHfExportModel(NamedTuple):
             collect_method="min-mean",
             model_dtype="BFloat16",
         ),
-        pytest.param(
-            DiffuserHfExportModel(
-                name="sdxl-1.0",
-                path=SDXL_PATH,
-                dtype="Half",
-                format_type="fp8",
-                quant_algo="max",
-                collect_method="default",
-            ),
-            marks=minimum_sm(89),
-        ),
         pytest.param(
             DiffuserHfExportModel(
                 name="flux-schnell",
@@ -115,22 +101,12 @@ class DiffuserHfExportModel(NamedTuple):
         ),
     ],
     ids=[
-        "sdxl_1.0_int8_smoothquant_min_mean",
         "flux_schnell_int8_smoothquant_min_mean",
-        "sdxl_1.0_fp8_max_default",
         "flux_schnell_fp4_max_default",
     ],
 )
 def test_diffusers_hf_ckpt_export(model: DiffuserHfExportModel, tmp_path: Path) -> None:
-    hf_ckpt_dir = model.quantize_and_export_hf(tmp_path)
-
-    assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}"
-
-    config_files = list(hf_ckpt_dir.rglob("config.json"))
-    assert len(config_files) > 0, f"No config.json found in {hf_ckpt_dir}"
-
-    weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin"))
-    assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}"
+    assert_hf_ckpt_exported(model.quantize_and_export_hf(tmp_path))
 
 
 class QwenHfExportModel(NamedTuple):
@@ -163,6 +139,8 @@ class QwenHfExportModel(NamedTuple):
             "1",
             "--n-steps",
             "2",
+            "--prompts-file",
+            str(CALIB_PROMPTS_FILE),
             "--hf-ckpt-dir",
             str(hf_ckpt_dir),
         ]
@@ -306,84 +284,3 @@ def test_qwen_image_hf_ckpt_export(
         assert not any(k.endswith(".svdquant_lora_a") for k in keys)
         if qwen_model.format_type == "fp4":
             assert any(k.endswith(".weight_scale_2") for k in keys)
-
-
-class Wan22HfExportModel(NamedTuple):
-    model: str
-    backbone: str | None
-    format_type: str
-    quant_algo: str
-    collect_method: str
-
-    def _suffix(self) -> str:
-        stem = self.model.replace("wan2.2-t2v-", "")
-        parts = [stem, *([self.backbone] if self.backbone else []), self.format_type]
-        return "_".join(parts)
-
-    def quantize_and_export_hf(self, tiny_wan22_path: str, tmp_path: Path) -> Path:
-        hf_ckpt_dir = tmp_path / f"wan22_{self._suffix()}_hf_ckpt"
-        cmd_args = [
-            "python",
-            "quantize.py",
-            "--model",
-            self.model,
-            "--override-model-path",
-            tiny_wan22_path,
-            "--format",
-            self.format_type,
-            "--quant-algo",
-            self.quant_algo,
-            "--collect-method",
-            self.collect_method,
-            "--model-dtype",
-            "BFloat16",
-            "--trt-high-precision-dtype",
-            "BFloat16",
-            "--calib-size",
-            "2",
-            "--batch-size",
-            "1",
-            "--n-steps",
-            "2",
-            # Tiny video dims — override MODEL_DEFAULTS for fast CI.
-            "--extra-param",
-            "height=16",
-            "--extra-param",
-            "width=16",
-            "--extra-param",
-            "num_frames=5",
-            "--hf-ckpt-dir",
-            str(hf_ckpt_dir),
-        ]
-        if self.backbone is not None:
-            cmd_args.extend(["--backbone", self.backbone])
-        run_example_command(cmd_args, "diffusers/quantization")
-        return hf_ckpt_dir
-
-
-@pytest.mark.parametrize(
-    "wan_model",
-    [
-        Wan22HfExportModel("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"),
-        pytest.param(
-            Wan22HfExportModel("wan2.2-t2v-14b", None, "fp8", "max", "default"),
-            marks=minimum_sm(89),
-        ),
-    ],
-    ids=[
-        "wan22_14b_transformer_int8_smoothquant",
-        "wan22_14b_transformer_fp8_max",
-    ],
-)
-def test_wan22_hf_ckpt_export(
-    wan_model: Wan22HfExportModel, tiny_wan22_path: str, tmp_path: Path
-) -> None:
-    hf_ckpt_dir = wan_model.quantize_and_export_hf(tiny_wan22_path, tmp_path)
-
-    assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}"
-
-    config_files = list(hf_ckpt_dir.rglob("config.json"))
-    assert len(config_files) > 0, f"No config.json found in {hf_ckpt_dir}"
-
-    weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin"))
-    assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}"