From aeda394db7d2f8a62475f2114da3b92e9588c42f Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Wed, 19 Aug 2026 10:32:03 +0530 Subject: [PATCH] [CI] Allow single file tests to run as fast tests (#14324) * update * update * update * update * [FA2] Use `kernels-community/aiter-flash-attn-ck` Hub kernel for ROCm FlashAttention (drop `aiter` dependency) (#14436) Use Hub kernel Co-authored-by: Sayak Paul * add HF_TOKEN to pr workflow * fix token * clean up * Apply suggestions from code review Co-authored-by: Sayak Paul --------- Co-authored-by: Sayak Paul Co-authored-by: BADAOUI Abdennacer <106801897+Abdennacer-Badaoui@users.noreply.github.com> --- .github/workflows/pr_tests.yml | 6 + src/diffusers/loaders/single_file_utils.py | 18 +- .../test_models_autoencoder_dc.py | 18 +- .../test_models_autoencoder_kl.py | 24 +- .../test_models_autoencoder_ltx_video.py | 22 +- .../test_models_autoencoder_wan.py | 22 +- tests/models/testing_utils/single_file.py | 255 ++++++++---------- .../test_models_transformer_chroma.py | 15 ++ .../test_models_transformer_flux.py | 4 - .../test_models_transformer_flux2.py | 15 ++ .../test_models_transformer_hidream.py | 15 ++ .../test_models_transformer_ltx.py | 15 ++ .../test_models_transformer_lumina2.py | 15 ++ .../test_models_transformer_mochi.py | 15 ++ .../test_models_transformer_sana.py | 15 ++ .../test_models_transformer_sd3.py | 15 ++ .../test_models_transformer_wan.py | 33 +++ .../test_models_transformer_wan_vace.py | 15 ++ .../test_models_transformer_z_image.py | 15 ++ .../unets/test_models_unet_2d_condition.py | 15 ++ 20 files changed, 408 insertions(+), 159 deletions(-) diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index ebb19a81b7..9885fb76d2 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -114,6 +114,9 @@ jobs: run: shell: bash + env: + HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} + steps: - name: Checkout diffusers uses: actions/checkout@v6 @@ -242,6 +245,9 @@ jobs: run: shell: bash + env: + HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }} + steps: - name: Checkout diffusers uses: actions/checkout@v6 diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index b5c6846ba1..ac707947d4 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -149,6 +149,10 @@ CHECKPOINT_KEY_NAMES = { "net.pos_embedder.dim_spatial_range", ], "flux2": ["model.diffusion_model.single_stream_modulation.lin.weight", "single_stream_modulation.lin.weight"], + "chroma": [ + "model.diffusion_model.distilled_guidance_layer.in_proj.bias", + "distilled_guidance_layer.in_proj.bias", + ], "ltx2": [ "model.diffusion_model.av_ca_a2v_gate_adaln_single.emb.timestep_embedder.linear_1.weight", "vae.per_channel_statistics.mean-of-means", @@ -204,6 +208,7 @@ DIFFUSERS_DEFAULT_PIPELINE_PATHS = { "flux-depth": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-Depth-dev"}, "flux-schnell": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-schnell"}, "flux-2-dev": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.2-dev"}, + "chroma": {"pretrained_model_name_or_path": "lodestones/Chroma1-HD"}, "ltx-video": {"pretrained_model_name_or_path": "diffusers/LTX-Video-0.9.0"}, "ltx-video-0.9.1": {"pretrained_model_name_or_path": "diffusers/LTX-Video-0.9.1"}, "ltx-video-0.9.5": {"pretrained_model_name_or_path": "Lightricks/LTX-Video-0.9.5"}, @@ -683,6 +688,9 @@ def infer_diffusers_model_type(checkpoint): elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux2"]): model_type = "flux-2-dev" + elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["chroma"]): + model_type = "chroma" + elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux"]): if any( g in checkpoint for g in ["guidance_in.in_layer.bias", "model.diffusion_model.guidance_in.in_layer.bias"] @@ -754,17 +762,19 @@ def infer_diffusers_model_type(checkpoint): elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["wan"]): if "model.diffusion_model.patch_embedding.weight" in checkpoint: - target_key = "model.diffusion_model.patch_embedding.weight" + prefix = "model.diffusion_model." else: - target_key = "patch_embedding.weight" + prefix = "" - if CHECKPOINT_KEY_NAMES["wan_vace"] in checkpoint: + target_key = f"{prefix}patch_embedding.weight" + + if f"{prefix}{CHECKPOINT_KEY_NAMES['wan_vace']}" in checkpoint: if checkpoint[target_key].shape[0] == 1536: model_type = "wan-vace-1.3B" elif checkpoint[target_key].shape[0] == 5120: model_type = "wan-vace-14B" - if CHECKPOINT_KEY_NAMES["wan_animate"] in checkpoint: + elif f"{prefix}{CHECKPOINT_KEY_NAMES['wan_animate']}" in checkpoint: model_type = "wan-animate-14B" elif checkpoint[target_key].shape[0] == 1536: diff --git a/tests/models/autoencoders/test_models_autoencoder_dc.py b/tests/models/autoencoders/test_models_autoencoder_dc.py index a5ce0a975a..c174362302 100644 --- a/tests/models/autoencoders/test_models_autoencoder_dc.py +++ b/tests/models/autoencoders/test_models_autoencoder_dc.py @@ -20,7 +20,13 @@ from diffusers import AutoencoderDC from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import IS_GITHUB_ACTIONS, enable_full_determinism, torch_device -from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + SingleFileTesterMixin, + TrainingTesterMixin, +) from .testing_utils import NewAutoencoderTesterMixin @@ -102,3 +108,13 @@ class TestAutoencoderDCMemory(AutoencoderDCTesterConfig, MemoryTesterMixin): class TestAutoencoderDCSlicingTiling(AutoencoderDCTesterConfig, NewAutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderDC.""" + + +class TestAutoencoderDCSingleFile(AutoencoderDCTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 0e209b5219..2ab286ffb8 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -35,7 +35,13 @@ from ...testing_utils import ( torch_all_close, torch_device, ) -from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + SingleFileTesterMixin, + TrainingTesterMixin, +) from .testing_utils import NewAutoencoderTesterMixin @@ -428,3 +434,19 @@ class AutoencoderKLIntegrationTests: tolerance = 3e-3 if torch_device != "mps" else 1e-2 assert torch_all_close(output_slice, expected_output_slice, atol=tolerance) + + +class TestAutoencoderKLSingleFile(AutoencoderKLTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors" + + @property + def pretrained_model_name_or_path(self): + # `from_single_file` resolves a bare VAE checkpoint against the SD 1.5 repo, so compare against that + # rather than against `stabilityai/sd-vae-ft-mse`, whose own config differs (sample_size 256 vs 512). + return "stable-diffusion-v1-5/stable-diffusion-v1-5" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "vae"} diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py index d93c1e5667..b8d992c63d 100644 --- a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py @@ -20,7 +20,13 @@ from diffusers import AutoencoderKLLTXVideo from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device -from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + SingleFileTesterMixin, + TrainingTesterMixin, +) from .testing_utils import NewAutoencoderTesterMixin @@ -179,3 +185,17 @@ class TestAutoencoderKLLTXVideo091Training(AutoencoderKLLTXVideo091TesterConfig, class TestAutoencoderKLLTXVideo091Memory(AutoencoderKLLTXVideo091TesterConfig, MemoryTesterMixin): """Memory optimization tests for AutoencoderKLLTXVideo (0.9.1 config).""" + + +class TestAutoencoderKLLTXVideoSingleFile(AutoencoderKLLTXVideo090TesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "diffusers/LTX-Video-0.9.0" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "vae"} diff --git a/tests/models/autoencoders/test_models_autoencoder_wan.py b/tests/models/autoencoders/test_models_autoencoder_wan.py index 89e5c58e4a..ee67cc87a9 100644 --- a/tests/models/autoencoders/test_models_autoencoder_wan.py +++ b/tests/models/autoencoders/test_models_autoencoder_wan.py @@ -20,7 +20,13 @@ from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device -from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + SingleFileTesterMixin, + TrainingTesterMixin, +) from .testing_utils import NewAutoencoderTesterMixin @@ -86,3 +92,17 @@ class TestAutoencoderKLWanMemory(AutoencoderKLWanTesterConfig, MemoryTesterMixin class TestAutoencoderKLWanSlicingTiling(AutoencoderKLWanTesterConfig, NewAutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLWan.""" + + +class TestAutoencoderKLWanSingleFile(AutoencoderKLWanTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "vae"} diff --git a/tests/models/testing_utils/single_file.py b/tests/models/testing_utils/single_file.py index 6e9443aa08..0d19baf656 100644 --- a/tests/models/testing_utils/single_file.py +++ b/tests/models/testing_utils/single_file.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import gc +import huggingface_hub +import pytest import torch -from huggingface_hub import hf_hub_download, snapshot_download +from huggingface_hub import HfApi +from safetensors.torch import _getdtype from diffusers.loaders.single_file_utils import _extract_repo_id_and_weights_name @@ -24,40 +28,30 @@ from ...testing_utils import ( backend_empty_cache, is_single_file, nightly, - require_torch_accelerator, torch_device, ) from .common import check_device_map_is_respected -def download_single_file_checkpoint(pretrained_model_name_or_path, filename, tmpdir): - """Download a single file checkpoint from the Hub to a temporary directory.""" - path = hf_hub_download(pretrained_model_name_or_path, filename=filename, local_dir=tmpdir) - return path +# Config entries that legitimately differ between a checkpoint that's loaded with `from_pretrained()` and the one loaded with `from_single_file()`. +PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] -def download_diffusers_config(pretrained_model_name_or_path, tmpdir): - """Download diffusers config files (excluding weights) from a repository.""" - path = snapshot_download( - pretrained_model_name_or_path, - ignore_patterns=[ - "**/*.ckpt", - "*.ckpt", - "**/*.bin", - "*.bin", - "**/*.pt", - "*.pt", - "**/*.safetensors", - "*.safetensors", - ], - allow_patterns=["**/*.json", "*.json", "*.txt", "**/*.txt"], - local_dir=tmpdir, - ) - return path +@functools.lru_cache(maxsize=None) +def fetch_checkpoint_metadata(ckpt_path): + """ + Fetch a single file checkpoint's keys, shapes and dtypes without downloading its weights. + + Only the safetensors header crosses the network, and the result is cached for the session because every test + in a class asks for the same one. Returns `None` for checkpoints that are not safetensors. + """ + pretrained_model_name_or_path, weight_name = _extract_repo_id_and_weights_name(ckpt_path) + if not weight_name.endswith(".safetensors"): + return None + + return HfApi().parse_safetensors_file_metadata(pretrained_model_name_or_path, weight_name) -@nightly -@require_torch_accelerator @is_single_file class SingleFileTesterMixin: """ @@ -68,7 +62,6 @@ class SingleFileTesterMixin: Optional properties: - torch_dtype: torch dtype to use for testing (default: None) - - alternate_ckpt_paths: List of alternate checkpoint paths for variant testing (default: None) Expected from config mixin: - model_class: The model class to test @@ -93,11 +86,6 @@ class SingleFileTesterMixin: """torch dtype to use for single file testing.""" return None - @property - def alternate_ckpt_paths(self) -> list[str] | None: - """List of alternate checkpoint paths for variant testing.""" - return None - def setup_method(self): gc.collect() backend_empty_cache(torch_device) @@ -106,47 +94,112 @@ class SingleFileTesterMixin: gc.collect() backend_empty_cache(torch_device) + def get_dummy_single_file_state_dict(self): + """ + A stand-in for the single file checkpoint's state dict. + + Keys, shapes and dtypes match the real checkpoint exactly, while the tensors themselves are on meta and + hold no data. Pass it to `from_single_file` with `device="meta"` to exercise the loading path without + materializing a weight. The conversion functions consume the dict they are handed, so build a fresh one + for every load. + """ + metadata = fetch_checkpoint_metadata(self.ckpt_path) + if metadata is None: + pytest.skip(f"{self.ckpt_path} is not a safetensors checkpoint") + + return { + name: torch.empty(tuple(tensor.shape), dtype=_getdtype(tensor.dtype), device="meta") + for name, tensor in metadata.tensors.items() + } + def test_single_file_model_config(self): - pretrained_kwargs = {"device_map": "auto", **self.pretrained_model_kwargs} - single_file_kwargs = {"device_map": "auto"} + # An empty model supplies the registered defaults (`out_channels` and friends) that a bare config dict + # would be missing. + with torch.device("meta"): + model = self.model_class.from_config(self.pretrained_model_name_or_path, **self.pretrained_model_kwargs) - if self.torch_dtype: - pretrained_kwargs["torch_dtype"] = self.torch_dtype - single_file_kwargs["torch_dtype"] = self.torch_dtype + config = model.config + model_single_file = self.model_class.from_single_file(self.get_dummy_single_file_state_dict(), device="meta") - model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) - model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) - - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] for param_name, param_value in model_single_file.config.items(): if param_name in PARAMS_TO_IGNORE: continue - assert model.config[param_name] == param_value, ( + assert config[param_name] == param_value, ( f"{param_name} differs between pretrained loading and single file loading: " - f"pretrained={model.config[param_name]}, single_file={param_value}" + f"pretrained={config[param_name]}, single_file={param_value}" ) + def test_single_file_loading_local_files_only(self, monkeypatch): + state_dict = self.get_dummy_single_file_state_dict() + + # The checkpoint is already in memory, so the config is the only thing left to resolve. Fetch it into the + # cache the way a previous run would have, and keep it as the reference to compare against. + config = self.model_class.load_config(self.pretrained_model_name_or_path, **self.pretrained_model_kwargs) + + # Cut the Hub off, so anything the load does not find locally raises instead of quietly downloading. + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) + + model_single_file = self.model_class.from_single_file(state_dict, local_files_only=True, device="meta") + + # Resolving from the cache has to land on the same config as resolving from the Hub. + for param_name, param_value in config.items(): + if param_name in PARAMS_TO_IGNORE: + continue + assert model_single_file.config[param_name] == param_value, ( + f"{param_name} differs when loading with local_files_only=True: " + f"pretrained={param_value}, single_file={model_single_file.config[param_name]}" + ) + + def test_single_file_loading_with_diffusers_config(self): + model_single_file = self.model_class.from_single_file( + self.get_dummy_single_file_state_dict(), + config=self.pretrained_model_name_or_path, + device="meta", + **self.pretrained_model_kwargs, + ) + + # An empty model supplies the registered defaults (`out_channels` and friends) that a bare config dict + # would be missing. + with torch.device("meta"): + model = self.model_class.from_config(self.pretrained_model_name_or_path, **self.pretrained_model_kwargs) + + config = model.config + + # Compare configs + for param_name, param_value in model_single_file.config.items(): + if param_name in PARAMS_TO_IGNORE: + continue + assert config[param_name] == param_value, ( + f"{param_name} differs: pretrained={config[param_name]}, single_file={param_value}" + ) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) + def test_single_file_loading_dtype(self, dtype): + model_single_file = self.model_class.from_single_file( + self.get_dummy_single_file_state_dict(), device="meta", dtype=dtype + ) + for name, param in model_single_file.named_parameters(): + assert param.dtype == dtype, f"Parameter {name} should be {dtype} but got {param.dtype}" + + @nightly def test_single_file_model_parameters(self): - pretrained_kwargs = {"device_map": "auto", **self.pretrained_model_kwargs} - single_file_kwargs = {"device_map": "auto"} + pretrained_kwargs = {**self.pretrained_model_kwargs} + single_file_kwargs = {} if self.torch_dtype: - pretrained_kwargs["torch_dtype"] = self.torch_dtype - single_file_kwargs["torch_dtype"] = self.torch_dtype + pretrained_kwargs["dtype"] = self.torch_dtype + single_file_kwargs["dtype"] = self.torch_dtype - # Load pretrained model, get state dict on CPU, then free GPU memory + # Both models load onto the CPU, so keep only the state dicts and drop the models themselves. model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) - state_dict = {k: v.cpu() for k, v in model.state_dict().items()} + state_dict = model.state_dict() del model gc.collect() - backend_empty_cache(torch_device) - # Load single file model, get state dict on CPU model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) - state_dict_single_file = {k: v.cpu() for k, v in model_single_file.state_dict().items()} + state_dict_single_file = model_single_file.state_dict() del model_single_file gc.collect() - backend_empty_cache(torch_device) assert set(state_dict.keys()) == set(state_dict_single_file.keys()), ( "Model parameters keys differ between pretrained and single file loading. " @@ -165,104 +218,12 @@ class SingleFileTesterMixin: assert torch.equal(param, param_single_file), f"Parameter values differ for {key}" - def test_single_file_loading_local_files_only(self, tmp_path): - single_file_kwargs = {} - - if self.torch_dtype: - single_file_kwargs["torch_dtype"] = self.torch_dtype - - pretrained_model_name_or_path, weight_name = _extract_repo_id_and_weights_name(self.ckpt_path) - local_ckpt_path = download_single_file_checkpoint(pretrained_model_name_or_path, weight_name, str(tmp_path)) - - model_single_file = self.model_class.from_single_file( - local_ckpt_path, local_files_only=True, **single_file_kwargs - ) - - assert model_single_file is not None, "Failed to load model with local_files_only=True" - - def test_single_file_loading_with_diffusers_config(self): - single_file_kwargs = {} - - if self.torch_dtype: - single_file_kwargs["torch_dtype"] = self.torch_dtype - single_file_kwargs.update(self.pretrained_model_kwargs) - - # Load with config parameter - model_single_file = self.model_class.from_single_file( - self.ckpt_path, config=self.pretrained_model_name_or_path, **single_file_kwargs - ) - - # Load pretrained for comparison - pretrained_kwargs = {**self.pretrained_model_kwargs} - if self.torch_dtype: - pretrained_kwargs["torch_dtype"] = self.torch_dtype - - model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) - - # Compare configs - PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] - for param_name, param_value in model_single_file.config.items(): - if param_name in PARAMS_TO_IGNORE: - continue - assert model.config[param_name] == param_value, ( - f"{param_name} differs: pretrained={model.config[param_name]}, single_file={param_value}" - ) - - def test_single_file_loading_with_diffusers_config_local_files_only(self, tmp_path): - single_file_kwargs = {} - - if self.torch_dtype: - single_file_kwargs["torch_dtype"] = self.torch_dtype - single_file_kwargs.update(self.pretrained_model_kwargs) - - pretrained_model_name_or_path, weight_name = _extract_repo_id_and_weights_name(self.ckpt_path) - local_ckpt_path = download_single_file_checkpoint(pretrained_model_name_or_path, weight_name, str(tmp_path)) - local_diffusers_config = download_diffusers_config(self.pretrained_model_name_or_path, str(tmp_path)) - - model_single_file = self.model_class.from_single_file( - local_ckpt_path, config=local_diffusers_config, local_files_only=True, **single_file_kwargs - ) - - assert model_single_file is not None, "Failed to load model with config and local_files_only=True" - - def test_single_file_loading_dtype(self): - for dtype in [torch.float32, torch.float16]: - if torch_device == "mps" and dtype == torch.bfloat16: - continue - - model_single_file = self.model_class.from_single_file(self.ckpt_path, torch_dtype=dtype) - - assert model_single_file.dtype == dtype, f"Expected dtype {dtype}, got {model_single_file.dtype}" - - # Cleanup - del model_single_file - gc.collect() - backend_empty_cache(torch_device) - - def test_checkpoint_variant_loading(self): - if not self.alternate_ckpt_paths: - return - - for ckpt_path in self.alternate_ckpt_paths: - backend_empty_cache(torch_device) - - single_file_kwargs = {} - if self.torch_dtype: - single_file_kwargs["torch_dtype"] = self.torch_dtype - - model = self.model_class.from_single_file(ckpt_path, **single_file_kwargs) - - assert model is not None, f"Failed to load checkpoint from {ckpt_path}" - - del model - gc.collect() - backend_empty_cache(torch_device) - + @nightly def test_single_file_loading_with_device_map(self): single_file_kwargs = {"device_map": "auto"} if self.torch_dtype: - single_file_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs["dtype"] = self.torch_dtype model = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) diff --git a/tests/models/transformers/test_models_transformer_chroma.py b/tests/models/transformers/test_models_transformer_chroma.py index d63aee32cb..dc300fbbe7 100644 --- a/tests/models/transformers/test_models_transformer_chroma.py +++ b/tests/models/transformers/test_models_transformer_chroma.py @@ -24,6 +24,7 @@ from ..testing_utils import ( LoraHotSwappingForModelTesterMixin, LoraTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TrainingTesterMixin, ) @@ -155,3 +156,17 @@ class TestChromaTransformerLoRAHotSwap(ChromaTransformerTesterConfig, LoraHotSwa "txt_ids": randn_tensor((sequence_length, num_image_channels), device=torch_device), "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size), } + + +class TestChromaTransformer2DSingleFile(ChromaTransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/lodestones/Chroma1-HD/blob/main/Chroma1-HD.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "lodestones/Chroma1-HD" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 1add8f81f6..eddce6ec19 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -340,10 +340,6 @@ class TestFluxSingleFile(FluxTransformerTesterConfig, SingleFileTesterMixin): def ckpt_path(self): return "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors" - @property - def alternate_ckpt_paths(self): - return ["https://huggingface.co/Comfy-Org/flux1-dev/blob/main/flux1-dev-fp8.safetensors"] - @property def pretrained_model_name_or_path(self): return "black-forest-labs/FLUX.1-dev" diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index b279b57dcc..8a00b34b89 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -36,6 +36,7 @@ from ..testing_utils import ( LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchAoCompileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, @@ -654,3 +655,17 @@ class TestFlux2TransformerKVCache(Flux2TransformerKVCacheTesterConfig): output = model(**base_config.get_dummy_inputs()) assert output.kv_cache is None + + +class TestFlux2Transformer2DSingleFile(Flux2TransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/black-forest-labs/FLUX.2-dev/blob/main/flux2-dev.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "black-forest-labs/FLUX.2-dev" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_hidream.py b/tests/models/transformers/test_models_transformer_hidream.py index c789850c01..6489d82f50 100644 --- a/tests/models/transformers/test_models_transformer_hidream.py +++ b/tests/models/transformers/test_models_transformer_hidream.py @@ -22,6 +22,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( BaseModelTesterConfig, ModelTesterMixin, + SingleFileTesterMixin, TrainingTesterMixin, ) @@ -106,3 +107,17 @@ class TestHiDreamTransformerTraining(HiDreamTransformerTesterConfig, TrainingTes def test_gradient_checkpointing_is_applied(self): expected_set = {"HiDreamImageTransformer2DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestHiDreamImageTransformer2DSingleFile(HiDreamTransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/HiDream-I1_ComfyUI/blob/main/split_files/diffusion_models/hidream_i1_full_fp16.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "HiDream-ai/HiDream-I1-Dev" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_ltx.py b/tests/models/transformers/test_models_transformer_ltx.py index 4839703063..685734475b 100644 --- a/tests/models/transformers/test_models_transformer_ltx.py +++ b/tests/models/transformers/test_models_transformer_ltx.py @@ -22,6 +22,7 @@ from ..testing_utils import ( BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -116,3 +117,17 @@ class TestLTXTransformerCompile(LTXTransformerTesterConfig, TorchCompileTesterMi # TODO: Add pretrained_model_name_or_path once a tiny LTX model is available on the Hub # class TestLTXTransformerTorchAo(LTXTransformerTesterConfig, TorchAoTesterMixin): # """TorchAo quantization tests for LTX Video Transformer.""" + + +class TestLTXVideoTransformer3DSingleFile(LTXTransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "diffusers/LTX-Video-0.9.0" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_lumina2.py b/tests/models/transformers/test_models_transformer_lumina2.py index 7e50902392..1f54fb591e 100644 --- a/tests/models/transformers/test_models_transformer_lumina2.py +++ b/tests/models/transformers/test_models_transformer_lumina2.py @@ -22,6 +22,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( BaseModelTesterConfig, ModelTesterMixin, + SingleFileTesterMixin, TrainingTesterMixin, ) @@ -95,3 +96,17 @@ class TestLumina2TransformerTraining(Lumina2TransformerTesterConfig, TrainingTes def test_gradient_checkpointing_is_applied(self): expected_set = {"Lumina2Transformer2DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestLumina2Transformer2DSingleFile(Lumina2TransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/blob/main/split_files/diffusion_models/lumina_2_model_bf16.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "Alpha-VLLM/Lumina-Image-2.0" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_mochi.py b/tests/models/transformers/test_models_transformer_mochi.py index 8530dd98cf..2c282a4138 100644 --- a/tests/models/transformers/test_models_transformer_mochi.py +++ b/tests/models/transformers/test_models_transformer_mochi.py @@ -22,6 +22,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( BaseModelTesterConfig, ModelTesterMixin, + SingleFileTesterMixin, TrainingTesterMixin, ) @@ -98,3 +99,17 @@ class TestMochiTransformerTraining(MochiTransformerTesterConfig, TrainingTesterM def test_gradient_checkpointing_is_applied(self): expected_set = {"MochiTransformer3DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestMochiTransformer3DSingleFile(MochiTransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/mochi_preview_repackaged/blob/main/split_files/diffusion_models/mochi_preview_bf16.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "genmo/mochi-1-preview" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_sana.py b/tests/models/transformers/test_models_transformer_sana.py index 41d581c290..c162dabd69 100644 --- a/tests/models/transformers/test_models_transformer_sana.py +++ b/tests/models/transformers/test_models_transformer_sana.py @@ -23,6 +23,7 @@ from ..testing_utils import ( BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TrainingTesterMixin, ) @@ -113,3 +114,17 @@ class TestSanaTransformerTraining(SanaTransformerTesterConfig, TrainingTesterMix class TestSanaTransformerAttention(SanaTransformerTesterConfig, AttentionTesterMixin): pass + + +class TestSanaTransformer2DSingleFile(SanaTransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Efficient-Large-Model/Sana_1600M_1024px/blob/main/checkpoints/Sana_1600M_1024px.pth" + + @property + def pretrained_model_name_or_path(self): + return "Efficient-Large-Model/Sana_1600M_1024px_diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index 2f1d5fff6e..a54b8e2d3f 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -23,6 +23,7 @@ from ..testing_utils import ( BaseModelTesterConfig, BitsAndBytesTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, @@ -268,3 +269,17 @@ class TestSD35TransformerTorchAo(SD35TransformerQuantTesterConfig, TorchAoTester @property def torch_dtype(self): return torch.bfloat16 + + +class TestSD3Transformer2DSingleFile(SD3TransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/sd3_medium.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "stabilityai/stable-diffusion-3-medium-diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_wan.py b/tests/models/transformers/test_models_transformer_wan.py index c4c64eabcc..26210eb3d5 100644 --- a/tests/models/transformers/test_models_transformer_wan.py +++ b/tests/models/transformers/test_models_transformer_wan.py @@ -26,6 +26,7 @@ from ..testing_utils import ( GGUFTesterMixin, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, @@ -228,3 +229,35 @@ class TestWanTransformer3DGGUFCompile(WanTransformer3DTesterConfig, GGUFCompileT ), "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), } + + +class TestWanTransformer3DText2VideoSingleFile(WanTransformer3DTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + +class TestWanTransformer3DImage2VideoSingleFile(WanTransformer3DTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/diffusion_models/wan2.1_i2v_480p_14B_fp8_e4m3fn.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + @property + def torch_dtype(self): + return torch.float8_e4m3fn diff --git a/tests/models/transformers/test_models_transformer_wan_vace.py b/tests/models/transformers/test_models_transformer_wan_vace.py index ac980b1ea7..cc0a3e34e1 100644 --- a/tests/models/transformers/test_models_transformer_wan_vace.py +++ b/tests/models/transformers/test_models_transformer_wan_vace.py @@ -27,6 +27,7 @@ from ..testing_utils import ( GGUFTesterMixin, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, @@ -251,3 +252,17 @@ class TestWanVACETransformer3DGGUFCompile(WanVACETransformer3DTesterConfig, GGUF ), "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), } + + +class TestWanVACETransformer3DSingleFile(WanVACETransformer3DTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/diffusion_models/wan2.1_vace_1.3B_fp16.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "Wan-AI/Wan2.1-VACE-1.3B-diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 57e40611a2..695b54be4c 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -28,6 +28,7 @@ from ..testing_utils import ( LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -358,3 +359,17 @@ class TestZImageTransformerAutoRoundCompile(ZImageTransformerAutoRoundTesterConf output = output[0] if isinstance(output, (list, tuple)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" + + +class TestZImageTransformer2DSingleFile(ZImageTransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "Tongyi-MAI/Z-Image-Turbo" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} diff --git a/tests/models/unets/test_models_unet_2d_condition.py b/tests/models/unets/test_models_unet_2d_condition.py index fc50a0a461..f31d148caf 100644 --- a/tests/models/unets/test_models_unet_2d_condition.py +++ b/tests/models/unets/test_models_unet_2d_condition.py @@ -60,6 +60,7 @@ from ..testing_utils import ( LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -1443,3 +1444,17 @@ class UNet2DConditionModelIntegrationTests(unittest.TestCase): expected_output_slice = torch.tensor(expected_slice) assert torch_all_close(output_slice, expected_output_slice, atol=5e-3) + + +class TestUNet2DConditionSingleFile(UNet2DConditionTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors" + + @property + def pretrained_model_name_or_path(self): + return "stable-diffusion-v1-5/stable-diffusion-v1-5" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "unet"}