kanshan/diffusers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix loading sharded checkpoints when we have variants (#9061)

Browse Source 
* Fix loading sharded checkpoint when we have variant

* add test

* remote print

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
This commit is contained in:
Marc Sun
2024-08-07 01:38:44 +02:00
committed by GitHub
parent 926daa30f9
commit e4325606db
2 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/diffusers/models/modeling_utils.py
+2 -2
View File
@@ -773,7 +773,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
try:
accelerate.load_checkpoint_and_dispatch(
model,
model_file if not is_sharded else sharded_ckpt_cached_folder,
model_file if not is_sharded else index_file,
device_map,
max_memory=max_memory,
offload_folder=offload_folder,
@@ -803,7 +803,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
model._temp_convert_self_to_deprecated_attention_blocks()
accelerate.load_checkpoint_and_dispatch(
model,
model_file if not is_sharded else sharded_ckpt_cached_folder,
model_file if not is_sharded else index_file,
device_map,
max_memory=max_memory,
offload_folder=offload_folder,
tests/models/unets/test_models_unet_2d_condition.py
+12
View File
@@ -1121,6 +1121,18 @@ class UNet2DConditionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.Test
assert loaded_model
assert new_output.sample.shape == (4, 4, 16, 16)
@require_torch_gpu
def test_load_sharded_checkpoint_with_variant_from_hub(self):
_, inputs_dict = self.prepare_init_args_and_inputs_for_common()
loaded_model = self.model_class.from_pretrained(
"hf-internal-testing/unet2d-sharded-with-variant-dummy", variant="fp16"
)
loaded_model = loaded_model.to(torch_device)
new_output = loaded_model(**inputs_dict)
assert loaded_model
assert new_output.sample.shape == (4, 4, 16, 16)
@require_peft_backend
def test_lora(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()