[FEAT] DDUF format

src/diffusers/pipelines/pipeline_utils.py

@@ -193,6 +193,7 @@ def save_pretrained(
         variant: Optional[str] = None,
         max_shard_size: Optional[Union[int, str]] = None,
         push_to_hub: bool = False,
+        dduf_filename: Optional[Union[str, os.PathLike]] = None,
         **kwargs,
     ):
         """
@@ -301,9 +302,54 @@ def is_saveable_module(name, value):
 
             save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs)
 
+            if dduf_filename:
+                import shutil
+                import zipfile
+
+                def zipdir(dir_to_archive, zipf):
+                    "zip a directory"
+                    for root, dirs, files in os.walk(dir_to_archive):
+                        for file in files:
+                            file_path = os.path.join(root, file)
+                            arcname = os.path.join(
+                                os.path.basename(dir_to_archive), os.path.relpath(file_path, start=dir_to_archive)
+                            )
+                            zipf.write(file_path, arcname=arcname)
+
+                dduf_file_path = os.path.join(save_directory, dduf_filename)
+
+                if os.path.isdir(dduf_file_path):
+                    logger.warning(
+                        f"Removing the existing folder {dduf_file_path} so that we can save the DDUF archive."
+                    )
+                    shutil.rmtree(dduf_file_path)
+                if (
+                    os.path.exists(dduf_file_path)
+                    and os.path.isfile(dduf_file_path)
+                    and zipfile.is_zipfile(dduf_file_path)
+                ):
+                    # Open in append mode if the file exists
+                    mode = "a"
+                else:
+                    # Open in write mode to create it if it doesn't exist
+                    mode = "w"
+                with zipfile.ZipFile(dduf_file_path, mode=mode, compression=zipfile.ZIP_STORED) as zipf:
+                    dir_to_archive = os.path.join(save_directory, pipeline_component_name)
+                    if os.path.isdir(dir_to_archive):
+                        zipdir(dir_to_archive, zipf)
+                        shutil.rmtree(dir_to_archive)
+
         # finally save the config
         self.save_config(save_directory)
 
+        if dduf_filename:
+            import zipfile
+
+            with zipfile.ZipFile(dduf_file_path, mode="a", compression=zipfile.ZIP_STORED) as zipf:
+                config_path = os.path.join(save_directory, self.config_name)
+                zipf.write(config_path, arcname=os.path.basename(config_path))
+                os.remove(config_path)
+
         if push_to_hub:
             # Create a new empty model card and eventually tag it
             model_card = load_or_create_model_card(repo_id, token=token, is_pipeline=True)
@@ -523,6 +569,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                     - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights
                       saved using
                     [`~DiffusionPipeline.save_pretrained`].
+                    - A path to a *directory* (for example `./my_pipeline_directory/`) containing a dduf archive or
+                      folder
             torch_dtype (`str` or `torch.dtype`, *optional*):
                 Override the default `torch.dtype` and load the model with another dtype. If "auto" is passed, the
                 dtype is automatically derived from the model's weights.
@@ -617,6 +665,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             variant (`str`, *optional*):
                 Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when
                 loading `from_flax`.
+            dduf (`str`, *optional*):
+                Load weights from the specified dduf archive or folder.
 
         <Tip>
 
@@ -666,6 +716,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         offload_state_dict = kwargs.pop("offload_state_dict", False)
         low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
         variant = kwargs.pop("variant", None)
+        dduf = kwargs.pop("dduf", None)
         use_safetensors = kwargs.pop("use_safetensors", None)
         use_onnx = kwargs.pop("use_onnx", None)
         load_connected_pipeline = kwargs.pop("load_connected_pipeline", False)
@@ -736,6 +787,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 custom_pipeline=custom_pipeline,
                 custom_revision=custom_revision,
                 variant=variant,
+                dduf=dduf,
                 load_connected_pipeline=load_connected_pipeline,
                 **kwargs,
             )
@@ -757,6 +809,25 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             )
             logger.warning(warn_msg)
 
+        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)
+                    # 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
+
         config_dict = cls.load_config(cached_folder)
 
         # pop out "_ignore_files" as it is only needed for download
@@ -1227,6 +1298,8 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
             variant (`str`, *optional*):
                 Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when
                 loading `from_flax`.
+            dduf(`str`, *optional*):
+                Load weights from the specified DDUF archive or folder.
             use_safetensors (`bool`, *optional*, defaults to `None`):
                 If set to `None`, the safetensors weights are downloaded if they're available **and** if the
                 safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors
@@ -1267,6 +1340,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
         use_onnx = kwargs.pop("use_onnx", None)
         load_connected_pipeline = kwargs.pop("load_connected_pipeline", False)
         trust_remote_code = kwargs.pop("trust_remote_code", False)
+        dduf = kwargs.pop("dduf", False)
 
         allow_pickle = False
         if use_safetensors is None:
@@ -1346,6 +1420,10 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
             allow_patterns += [f"{custom_pipeline}.py"] if f"{custom_pipeline}.py" in filenames else []
             # also allow downloading config.json files with the model
             allow_patterns += [os.path.join(k, "config.json") for k in model_folder_names]
+            # also allow downloading the dduf
+            # TODO: check that the file actually exist
+            if dduf is not None:
+                allow_patterns += [dduf]
 
             allow_patterns += [
                 SCHEDULER_CONFIG_NAME,