first draft

Author: SunMarc

src/diffusers/configuration_utils.py

@@ -347,6 +347,7 @@ def load_config(
         _ = kwargs.pop("mirror", None)
         subfolder = kwargs.pop("subfolder", None)
         user_agent = kwargs.pop("user_agent", {})
+        dduf_reader = kwargs.pop("dduf_reader", None)
 
         user_agent = {**user_agent, "file_type": "config"}
         user_agent = http_user_agent(user_agent)
@@ -358,8 +359,25 @@ def load_config(
                 "`self.config_name` is not defined. Note that one should not load a config from "
                 "`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`"
             )
-
-        if os.path.isfile(pretrained_model_name_or_path):
+        # Custom path for now
+        if dduf_reader:
+            if subfolder is not None:
+                if dduf_reader.has_file(os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)):
+                    config_file = os.path.join(subfolder, cls.config_name)
+                else:
+                    raise ValueError(
+                        f"We did not manage to find the file {os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)} in the archive. We only have the following files {dduf_reader.files}"
+                    )
+            elif dduf_reader.has_file(os.path.join(pretrained_model_name_or_path, cls.config_name)):
+                config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
+            else:
+                raise ValueError(
+                    f"We did not manage to find the file {os.path.join(pretrained_model_name_or_path, cls.config_name)} in the archive. We only have the following files {dduf_reader.files}"
+                )
+            print(f"File found: {config_file}")
+        elif not dduf_reader:
+            print("not dduf")
+        elif os.path.isfile(pretrained_model_name_or_path):
             config_file = pretrained_model_name_or_path
         elif os.path.isdir(pretrained_model_name_or_path):
             if subfolder is not None and os.path.isfile(
@@ -426,10 +444,8 @@ def load_config(
                     f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
                     f"containing a {cls.config_name} file"
                 )
-
         try:
-            # Load config dict
-            config_dict = cls._dict_from_json_file(config_file)
+            config_dict = cls._dict_from_json_file(config_file, dduf_reader=dduf_reader)
 
             commit_hash = extract_commit_hash(config_file)
         except (json.JSONDecodeError, UnicodeDecodeError):
@@ -552,9 +568,12 @@ def extract_init_dict(cls, config_dict, **kwargs):
         return init_dict, unused_kwargs, hidden_config_dict
 
     @classmethod
-    def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]):
-        with open(json_file, "r", encoding="utf-8") as reader:
-            text = reader.read()
+    def _dict_from_json_file(cls, json_file: Union[str, os.PathLike], dduf_reader=None):
+        if dduf_reader:
+            text = dduf_reader.read_file(json_file, encoding="utf-8")
+        else:
+            with open(json_file, "r", encoding="utf-8") as reader:
+                text = reader.read()
         return json.loads(text)
 
     def __repr__(self):

src/diffusers/models/model_loading_utils.py

@@ -128,7 +128,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):
+def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None, dduf_reader=None):
     """
     Reads a checkpoint file, returning properly formatted errors if they arise.
     """
@@ -138,8 +138,15 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[
         return checkpoint_file
     try:
         file_extension = os.path.basename(checkpoint_file).split(".")[-1]
+        if dduf_reader:
+            checkpoint_file = dduf_reader.read_file(checkpoint_file)
         if file_extension == SAFETENSORS_FILE_EXTENSION:
-            return safetensors.torch.load_file(checkpoint_file, device="cpu")
+            if dduf_reader:
+                # tensors are loaded on cpu
+                return safetensors.torch.load(checkpoint_file)
+            else:
+                return safetensors.torch.load_file(checkpoint_file, device="cpu")
+
         else:
             weights_only_kwarg = {"weights_only": True} if is_torch_version(">=", "1.13") else {}
             return torch.load(
@@ -272,6 +279,7 @@ def _fetch_index_file(
     revision,
     user_agent,
     commit_hash,
+    dduf_reader=None,
 ):
     if is_local:
         index_file = Path(
@@ -297,6 +305,7 @@ def _fetch_index_file(
                 subfolder=None,
                 user_agent=user_agent,
                 commit_hash=commit_hash,
+                dduf_reader=dduf_reader,
             )
             index_file = Path(index_file)
         except (EntryNotFoundError, EnvironmentError):

src/diffusers/models/modeling_utils.py

@@ -557,6 +557,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         variant = kwargs.pop("variant", None)
         use_safetensors = kwargs.pop("use_safetensors", None)
         quantization_config = kwargs.pop("quantization_config", None)
+        dduf_reader = kwargs.pop("dduf_reader", None)
 
         allow_pickle = False
         if use_safetensors is None:
@@ -649,6 +650,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             revision=revision,
             subfolder=subfolder,
             user_agent=user_agent,
+            dduf_reader=dduf_reader,
             **kwargs,
         )
         # no in-place modification of the original config.
@@ -724,6 +726,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             "revision": revision,
             "user_agent": user_agent,
             "commit_hash": commit_hash,
+            "dduf_reader": dduf_reader,
         }
         index_file = _fetch_index_file(**index_file_kwargs)
         # In case the index file was not found we still have to consider the legacy format.
@@ -759,7 +762,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
 
             model = load_flax_checkpoint_in_pytorch_model(model, model_file)
         else:
-            if is_sharded:
+            # in the case it is sharded, we have already the index
+            if is_sharded and not dduf_reader:
                 sharded_ckpt_cached_folder, sharded_metadata = _get_checkpoint_shard_files(
                     pretrained_model_name_or_path,
                     index_file,
@@ -790,6 +794,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                         subfolder=subfolder,
                         user_agent=user_agent,
                         commit_hash=commit_hash,
+                        dduf_reader=dduf_reader,
                     )
 
                 except IOError as e:
@@ -813,6 +818,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                     subfolder=subfolder,
                     user_agent=user_agent,
                     commit_hash=commit_hash,
+                    dduf_reader=dduf_reader,
                 )
 
             if low_cpu_mem_usage:
@@ -837,7 +843,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                     # TODO (sayakpaul,  SunMarc): remove this after model loading refactor
                     elif is_quant_method_bnb:
                         param_device = torch.cuda.current_device()
-                    state_dict = load_state_dict(model_file, variant=variant)
+                    state_dict = load_state_dict(model_file, variant=variant, dduf_reader=dduf_reader)
                     model._convert_deprecated_attention_blocks(state_dict)
 
                     # move the params from meta device to cpu
@@ -937,7 +943,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             else:
                 model = cls.from_config(config, **unused_kwargs)
 
-                state_dict = load_state_dict(model_file, variant=variant)
+                state_dict = load_state_dict(model_file, variant=variant, dduf_reader=dduf_reader)
                 model._convert_deprecated_attention_blocks(state_dict)
 
                 model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -627,6 +627,7 @@ def load_sub_model(
     low_cpu_mem_usage: bool,
     cached_folder: Union[str, os.PathLike],
     use_safetensors: bool,
+    dduf_reader,
 ):
     """Helper method to load the module `name` from `library_name` and `class_name`"""
 
@@ -721,7 +722,10 @@ def load_sub_model(
             loading_kwargs["low_cpu_mem_usage"] = False
 
     # check if the module is in a subdirectory
-    if os.path.isdir(os.path.join(cached_folder, name)):
+    if dduf_reader:
+        loading_kwargs["dduf_reader"] = dduf_reader
+        loaded_sub_model = load_method(name, **loading_kwargs)
+    elif os.path.isdir(os.path.join(cached_folder, name)):
         loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
     else:
         # else load from the root directory

src/diffusers/pipelines/pipeline_utils.py

@@ -50,6 +50,7 @@
     CONFIG_NAME,
     DEPRECATED_REVISION_ARGS,
     BaseOutput,
+    DDUFReader,
     PushToHubMixin,
     is_accelerate_available,
     is_accelerate_version,
@@ -343,7 +344,7 @@ def zipdir(dir_to_archive, zipf):
         self.save_config(save_directory)
 
         # Takes care of including the "model_index.json" inside the ZIP.
-        # TODO: Include a DDUF a metadata file. 
+        # TODO: Include a DDUF a metadata file.
         if dduf_filename:
             import zipfile
 
@@ -811,30 +812,14 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             )
             logger.warning(warn_msg)
 
+        dduf_reader = None
         if dduf:
-            import zipfile
-
             zip_file_path = os.path.join(cached_folder, dduf)
-            extract_to = os.path.join(cached_folder, f"{dduf}_extracted")
-            # if zip file, we need to extract the zipfile and remove it
-            if os.path.isfile(zip_file_path):
-                if zipfile.is_zipfile(zip_file_path):
-                    # with zipfile.ZipFile(zip_file_path, "r") as zipf:
-                    #     zipf.extractall(extract_to)
-                    with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
-                        file_list = zip_ref.infolist()
-                        for file_info in tqdm(file_list, desc="Extracting files"):
-                            zip_ref.extract(file_info, extract_to)
-                    # remove zip archive to free memory
-                    os.remove(zip_file_path)
-                    # rename folder to match the name of the dduf archive
-                    os.rename(extract_to, zip_file_path)
-                else:
-                    raise RuntimeError("The dduf path passed is not a zip archive")
-            # udapte cached folder location as the dduf content is in a seperate folder
-            cached_folder = zip_file_path
+            dduf_reader = DDUFReader(zip_file_path)
+            # The reader contains already all the files needed, no need to check it again
+            cached_folder = ""
 
-        config_dict = cls.load_config(cached_folder)
+        config_dict = cls.load_config(cached_folder, dduf_reader=dduf_reader)
 
         # pop out "_ignore_files" as it is only needed for download
         config_dict.pop("_ignore_files", None)
@@ -991,6 +976,7 @@ def load_module(name, value):
                     low_cpu_mem_usage=low_cpu_mem_usage,
                     cached_folder=cached_folder,
                     use_safetensors=use_safetensors,
+                    dduf_reader=dduf_reader,
                 )
                 logger.info(
                     f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}."

src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py

@@ -34,6 +34,7 @@
 from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
 from .modeling_stable_audio import StableAudioProjectionModel
 
+
 if is_torch_xla_available():
     import torch_xla.core.xla_model as xm
 
@@ -732,7 +733,7 @@ def __call__(
                     if callback is not None and i % callback_steps == 0:
                         step_idx = i // getattr(self.scheduler, "order", 1)
                         callback(step_idx, t, latents)
-                        
+
                 if XLA_AVAILABLE:
                     xm.mark_step()
 

src/diffusers/utils/__init__.py

@@ -35,6 +35,7 @@
     WEIGHTS_INDEX_NAME,
     WEIGHTS_NAME,
 )
+from .dduf import DDUFReader
 from .deprecation_utils import deprecate
 from .doc_utils import replace_example_docstring
 from .dynamic_modules_utils import get_class_from_dynamic_module

src/diffusers/utils/dduf.py

@@ -0,0 +1,41 @@
+import zipfile
+
+
+class DDUFReader:
+    def __init__(self, dduf_file):
+        self.dduf_file = dduf_file
+        self.files = []
+        self.post_init()
+
+    def post_init(self):
+        """
+        Check that the DDUF file is valid
+        """
+        if not zipfile.is_zipfile(self.dduf_file):
+            raise ValueError(f"The file '{self.dduf_file}' is not a valid ZIP archive.")
+
+        try:
+            with zipfile.ZipFile(self.dduf_file, "r") as zf:
+                # Check integrity and store file list
+                zf.testzip()  # Returns None if no corrupt files are found
+                self.files = zf.namelist()
+        except zipfile.BadZipFile:
+            raise ValueError(f"The file '{self.dduf_file}' is not a valid ZIP archive.")
+        except Exception as e:
+            raise RuntimeError(f"An error occurred while validating the ZIP file: {e}")
+
+    def has_file(self, file):
+        return file in self.files
+
+    def read_file(self, file_name, encoding=None):
+        """
+        Reads the content of a specific file in the ZIP archive without extracting.
+        """
+        if file_name not in self.files:
+            raise ValueError(f"{file_name} is not in the list of files {self.files}")
+        with zipfile.ZipFile(self.dduf_file, "r") as zf:
+            with zf.open(file_name) as file:
+                file = file.read()
+                if encoding is not None:
+                    file = file.decode(encoding)
+                return file

src/diffusers/utils/hub_utils.py

@@ -291,9 +291,20 @@ def _get_model_file(
     user_agent: Optional[Union[Dict, str]] = None,
     revision: Optional[str] = None,
     commit_hash: Optional[str] = None,
+    dduf_reader=None,
 ):
     pretrained_model_name_or_path = str(pretrained_model_name_or_path)
-    if os.path.isfile(pretrained_model_name_or_path):
+
+    if dduf_reader:
+        if dduf_reader.has_file(os.path.join(pretrained_model_name_or_path, weights_name)):
+            return os.path.join(pretrained_model_name_or_path, weights_name)
+        elif subfolder is not None and os.path.isfile(
+            os.path.join(pretrained_model_name_or_path, subfolder, weights_name)
+        ):
+            return os.path.join(pretrained_model_name_or_path, weights_name)
+        else:
+            raise EnvironmentError(f"Error no file named {weights_name} found in archive {dduf_reader.files}.")
+    elif os.path.isfile(pretrained_model_name_or_path):
         return pretrained_model_name_or_path
     elif os.path.isdir(pretrained_model_name_or_path):
         if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)):