Skip to content

Commit 710e18b

Browse files
DN6DN6
committed
update
1 parent 55f0b3d commit 710e18b

9 files changed

+96
-242
lines changed

‎tests/single_file/single_file_testing_utils.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import gc
12
import tempfile
23
from io import BytesIO
34

@@ -9,7 +10,9 @@
910
from diffusers.models.attention_processor import AttnProcessor
1011

1112
from ..testing_utils import (
13+
backend_empty_cache,
1214
numpy_cosine_similarity_distance,
15+
require_torch_accelerator,
1316
torch_device,
1417
)
1518

@@ -47,6 +50,76 @@ def download_diffusers_config(repo_id, tmpdir):
4750
return path
4851

4952

53+
@require_torch_accelerator
54+
class SingleFileModelTesterMixin:
55+
def setUp(self):
56+
super().setUp()
57+
gc.collect()
58+
backend_empty_cache(torch_device)
59+
60+
def tearDown(self):
61+
super().tearDown()
62+
gc.collect()
63+
backend_empty_cache(torch_device)
64+
65+
def test_single_file_model_config(self):
66+
pretrained_kwargs = {}
67+
single_file_kwargs = {}
68+
69+
if hasattr(self, "subfolder") and self.subfolder:
70+
pretrained_kwargs["subfolder"] = self.subfolder
71+
72+
if hasattr(self, "torch_dtype") and self.torch_dtype:
73+
pretrained_kwargs["torch_dtype"] = self.torch_dtype
74+
single_file_kwargs["torch_dtype"] = self.torch_dtype
75+
76+
model = self.model_class.from_pretrained(self.repo_id, **pretrained_kwargs)
77+
model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs)
78+
79+
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
80+
for param_name, param_value in model_single_file.config.items():
81+
if param_name in PARAMS_TO_IGNORE:
82+
continue
83+
assert model.config[param_name] == param_value, (
84+
f"{param_name} differs between pretrained loading and single file loading"
85+
)
86+
87+
def test_single_file_model_parameters(self):
88+
pretrained_kwargs = {}
89+
single_file_kwargs = {}
90+
91+
if hasattr(self, "subfolder") and self.subfolder:
92+
pretrained_kwargs["subfolder"] = self.subfolder
93+
94+
if hasattr(self, "torch_dtype") and self.torch_dtype:
95+
pretrained_kwargs["torch_dtype"] = self.torch_dtype
96+
single_file_kwargs["torch_dtype"] = self.torch_dtype
97+
98+
model = self.model_class.from_pretrained(self.repo_id, **pretrained_kwargs)
99+
model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs)
100+
101+
state_dict = model.state_dict()
102+
state_dict_single_file = model_single_file.state_dict()
103+
104+
assert set(state_dict.keys()) == set(state_dict_single_file.keys()), (
105+
"Model parameters keys differ between pretrained and single file loading"
106+
)
107+
108+
for key in state_dict.keys():
109+
param = state_dict[key]
110+
param_single_file = state_dict_single_file[key]
111+
112+
assert param.shape == param_single_file.shape, (
113+
f"Parameter shape mismatch for {key}: "
114+
f"pretrained {param.shape} vs single file {param_single_file.shape}"
115+
)
116+
117+
assert torch.allclose(param, param_single_file, rtol=1e-5, atol=1e-5), (
118+
f"Parameter values differ for {key}: "
119+
f"max difference {torch.max(torch.abs(param - param_single_file)).item()}"
120+
)
121+
122+
50123
class SDSingleFileTesterMixin:
51124
single_file_kwargs = {}
52125

‎tests/single_file/test_lumina2_transformer.py‎

Lines changed: 3 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -23,45 +23,23 @@
2323
from ..testing_utils import (
2424
backend_empty_cache,
2525
enable_full_determinism,
26-
require_torch_accelerator,
2726
torch_device,
2827
)
28+
from .single_file_testing_utils import SingleFileModelTesterMixin
2929

3030

3131
enable_full_determinism()
3232

3333

34-
@require_torch_accelerator
35-
class Lumina2Transformer2DModelSingleFileTests(unittest.TestCase):
34+
class Lumina2Transformer2DModelSingleFileTests(SingleFileModelTesterMixin, unittest.TestCase):
3635
model_class = Lumina2Transformer2DModel
3736
ckpt_path = "https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/blob/main/split_files/diffusion_models/lumina_2_model_bf16.safetensors"
3837
alternate_keys_ckpt_paths = [
3938
"https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/blob/main/split_files/diffusion_models/lumina_2_model_bf16.safetensors"
4039
]
4140

4241
repo_id = "Alpha-VLLM/Lumina-Image-2.0"
43-
44-
def setUp(self):
45-
super().setUp()
46-
gc.collect()
47-
backend_empty_cache(torch_device)
48-
49-
def tearDown(self):
50-
super().tearDown()
51-
gc.collect()
52-
backend_empty_cache(torch_device)
53-
54-
def test_single_file_components(self):
55-
model = self.model_class.from_pretrained(self.repo_id, subfolder="transformer")
56-
model_single_file = self.model_class.from_single_file(self.ckpt_path)
57-
58-
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
59-
for param_name, param_value in model_single_file.config.items():
60-
if param_name in PARAMS_TO_IGNORE:
61-
continue
62-
assert model.config[param_name] == param_value, (
63-
f"{param_name} differs between single file loading and pretrained loading"
64-
)
42+
subfolder = "transformer"
6543

6644
def test_checkpoint_loading(self):
6745
for ckpt_path in self.alternate_keys_ckpt_paths:

‎tests/single_file/test_model_autoencoder_dc_single_file.py‎

Lines changed: 2 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
import gc
1716
import unittest
1817

1918
import torch
@@ -23,38 +22,24 @@
2322
)
2423

2524
from ..testing_utils import (
26-
backend_empty_cache,
2725
enable_full_determinism,
2826
load_hf_numpy,
2927
numpy_cosine_similarity_distance,
30-
require_torch_accelerator,
31-
slow,
3228
torch_device,
3329
)
30+
from .single_file_testing_utils import SingleFileModelTesterMixin
3431

3532

3633
enable_full_determinism()
3734

3835

39-
@slow
40-
@require_torch_accelerator
41-
class AutoencoderDCSingleFileTests(unittest.TestCase):
36+
class AutoencoderDCSingleFileTests(SingleFileModelTesterMixin, unittest.TestCase):
4237
model_class = AutoencoderDC
4338
ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors"
4439
repo_id = "mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers"
4540
main_input_name = "sample"
4641
base_precision = 1e-2
4742

48-
def setUp(self):
49-
super().setUp()
50-
gc.collect()
51-
backend_empty_cache(torch_device)
52-
53-
def tearDown(self):
54-
super().tearDown()
55-
gc.collect()
56-
backend_empty_cache(torch_device)
57-
5843
def get_file_format(self, seed, shape):
5944
return f"gaussian_noise_s={seed}_shape={'_'.join([str(s) for s in shape])}.npy"
6045

@@ -80,18 +65,6 @@ def test_single_file_inference_same_as_pretrained(self):
8065

8166
assert numpy_cosine_similarity_distance(output_slice_1, output_slice_2) < 1e-4
8267

83-
def test_single_file_components(self):
84-
model = self.model_class.from_pretrained(self.repo_id)
85-
model_single_file = self.model_class.from_single_file(self.ckpt_path)
86-
87-
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
88-
for param_name, param_value in model_single_file.config.items():
89-
if param_name in PARAMS_TO_IGNORE:
90-
continue
91-
assert model.config[param_name] == param_value, (
92-
f"{param_name} differs between pretrained loading and single file loading"
93-
)
94-
9568
def test_single_file_in_type_variant_components(self):
9669
# `in` variant checkpoints require passing in a `config` parameter
9770
# in order to set the scaling factor correctly.

‎tests/single_file/test_model_controlnet_single_file.py‎

Lines changed: 2 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
import gc
1716
import unittest
1817

1918
import torch
@@ -23,46 +22,19 @@
2322
)
2423

2524
from ..testing_utils import (
26-
backend_empty_cache,
2725
enable_full_determinism,
28-
require_torch_accelerator,
29-
slow,
30-
torch_device,
3126
)
27+
from .single_file_testing_utils import SingleFileModelTesterMixin
3228

3329

3430
enable_full_determinism()
3531

3632

37-
@slow
38-
@require_torch_accelerator
39-
class ControlNetModelSingleFileTests(unittest.TestCase):
33+
class ControlNetModelSingleFileTests(SingleFileModelTesterMixin, unittest.TestCase):
4034
model_class = ControlNetModel
4135
ckpt_path = "https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_canny.pth"
4236
repo_id = "lllyasviel/control_v11p_sd15_canny"
4337

44-
def setUp(self):
45-
super().setUp()
46-
gc.collect()
47-
backend_empty_cache(torch_device)
48-
49-
def tearDown(self):
50-
super().tearDown()
51-
gc.collect()
52-
backend_empty_cache(torch_device)
53-
54-
def test_single_file_components(self):
55-
model = self.model_class.from_pretrained(self.repo_id)
56-
model_single_file = self.model_class.from_single_file(self.ckpt_path)
57-
58-
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
59-
for param_name, param_value in model_single_file.config.items():
60-
if param_name in PARAMS_TO_IGNORE:
61-
continue
62-
assert model.config[param_name] == param_value, (
63-
f"{param_name} differs between single file loading and pretrained loading"
64-
)
65-
6638
def test_single_file_arguments(self):
6739
model_default = self.model_class.from_single_file(self.ckpt_path)
6840

‎tests/single_file/test_model_flux_transformer_single_file.py‎

Lines changed: 3 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -23,43 +23,22 @@
2323
from ..testing_utils import (
2424
backend_empty_cache,
2525
enable_full_determinism,
26-
require_torch_accelerator,
2726
torch_device,
2827
)
28+
from .single_file_testing_utils import SingleFileModelTesterMixin
2929

3030

3131
enable_full_determinism()
3232

3333

34-
@require_torch_accelerator
35-
class FluxTransformer2DModelSingleFileTests(unittest.TestCase):
34+
class FluxTransformer2DModelSingleFileTests(SingleFileModelTesterMixin, unittest.TestCase):
3635
model_class = FluxTransformer2DModel
3736
ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
3837
alternate_keys_ckpt_paths = ["https://huggingface.co/Comfy-Org/flux1-dev/blob/main/flux1-dev-fp8.safetensors"]
3938

4039
repo_id = "black-forest-labs/FLUX.1-dev"
4140

42-
def setUp(self):
43-
super().setUp()
44-
gc.collect()
45-
backend_empty_cache(torch_device)
46-
47-
def tearDown(self):
48-
super().tearDown()
49-
gc.collect()
50-
backend_empty_cache(torch_device)
51-
52-
def test_single_file_components(self):
53-
model = self.model_class.from_pretrained(self.repo_id, subfolder="transformer")
54-
model_single_file = self.model_class.from_single_file(self.ckpt_path)
55-
56-
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
57-
for param_name, param_value in model_single_file.config.items():
58-
if param_name in PARAMS_TO_IGNORE:
59-
continue
60-
assert model.config[param_name] == param_value, (
61-
f"{param_name} differs between single file loading and pretrained loading"
62-
)
41+
subfolder = "transformer"
6342

6443
def test_checkpoint_loading(self):
6544
for ckpt_path in self.alternate_keys_ckpt_paths:

‎tests/single_file/test_model_vae_single_file.py‎

Lines changed: 2 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
import gc
1716
import unittest
1817

1918
import torch
@@ -23,22 +22,18 @@
2322
)
2423

2524
from ..testing_utils import (
26-
backend_empty_cache,
2725
enable_full_determinism,
2826
load_hf_numpy,
2927
numpy_cosine_similarity_distance,
30-
require_torch_accelerator,
31-
slow,
3228
torch_device,
3329
)
30+
from .single_file_testing_utils import SingleFileModelTesterMixin
3431

3532

3633
enable_full_determinism()
3734

3835

39-
@slow
40-
@require_torch_accelerator
41-
class AutoencoderKLSingleFileTests(unittest.TestCase):
36+
class AutoencoderKLSingleFileTests(SingleFileModelTesterMixin, unittest.TestCase):
4237
model_class = AutoencoderKL
4338
ckpt_path = (
4439
"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
@@ -47,16 +42,6 @@ class AutoencoderKLSingleFileTests(unittest.TestCase):
4742
main_input_name = "sample"
4843
base_precision = 1e-2
4944

50-
def setUp(self):
51-
super().setUp()
52-
gc.collect()
53-
backend_empty_cache(torch_device)
54-
55-
def tearDown(self):
56-
super().tearDown()
57-
gc.collect()
58-
backend_empty_cache(torch_device)
59-
6045
def get_file_format(self, seed, shape):
6146
return f"gaussian_noise_s={seed}_shape={'_'.join([str(s) for s in shape])}.npy"
6247

@@ -84,18 +69,6 @@ def test_single_file_inference_same_as_pretrained(self):
8469

8570
assert numpy_cosine_similarity_distance(output_slice_1, output_slice_2) < 1e-4
8671

87-
def test_single_file_components(self):
88-
model = self.model_class.from_pretrained(self.repo_id)
89-
model_single_file = self.model_class.from_single_file(self.ckpt_path, config=self.repo_id)
90-
91-
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
92-
for param_name, param_value in model_single_file.config.items():
93-
if param_name in PARAMS_TO_IGNORE:
94-
continue
95-
assert model.config[param_name] == param_value, (
96-
f"{param_name} differs between pretrained loading and single file loading"
97-
)
98-
9972
def test_single_file_arguments(self):
10073
model_default = self.model_class.from_single_file(self.ckpt_path, config=self.repo_id)
10174

0 commit comments

Comments
 (0)