typing

Author: SunMarc

src/diffusers/models/model_loading_utils.py

@@ -19,10 +19,11 @@
 import os
 from collections import OrderedDict
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import safetensors
 import torch
+from huggingface_hub import DDUFEntry
 from huggingface_hub.utils import EntryNotFoundError
 
 from ..quantizers.quantization_config import QuantizationMethod
@@ -128,7 +129,7 @@ def _fetch_remapped_cls_from_config(config, old_class):
         return old_class
 
 
-def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None, dduf_entries=None):
+def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None, dduf_entries: Optional[Dict[str, DDUFEntry]]=None):
     """
     Reads a checkpoint file, returning properly formatted errors if they arise.
     """

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -20,7 +20,7 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
-from huggingface_hub import ModelCard, model_info
+from huggingface_hub import DDUFEntry, ModelCard, model_info
 from huggingface_hub.utils import validate_hf_hub_args
 from packaging import version
 
@@ -626,7 +626,7 @@ def load_sub_model(
     low_cpu_mem_usage: bool,
     cached_folder: Union[str, os.PathLike],
     use_safetensors: bool,
-    dduf_entries,
+    dduf_entries: Optional[Dict[str, DDUFEntry]],
 ):
     """Helper method to load the module `name` from `library_name` and `class_name`"""
 

src/diffusers/utils/hub_utils.py

@@ -26,6 +26,7 @@
 from uuid import uuid4
 
 from huggingface_hub import (
+    DDUFEntry,
     ModelCard,
     ModelCardData,
     create_repo,
@@ -291,7 +292,7 @@ def _get_model_file(
     user_agent: Optional[Union[Dict, str]] = None,
     revision: Optional[str] = None,
     commit_hash: Optional[str] = None,
-    dduf_entries=None,
+    dduf_entries: Optional[Dict[str, DDUFEntry]]=None,
 ):
     pretrained_model_name_or_path = str(pretrained_model_name_or_path)