codekingpro/portable-devtools

1import json2import logging3import mmap4import os5import shutil6import zipfile7from collections.abc import Generator, Iterable8from contextlib import contextmanager9from dataclasses import dataclass, field10from pathlib import Path11from typing import Any12 13from ..errors import DDUFCorruptedFileError, DDUFExportError, DDUFInvalidEntryNameError14 15 16logger = logging.getLogger(__name__)17 18DDUF_ALLOWED_ENTRIES = {19    # Allowed file extensions in a DDUF file20    ".json",21    ".model",22    ".safetensors",23    ".txt",24}25 26DDUF_FOLDER_REQUIRED_ENTRIES = {27    # Each folder must contain at least one of these entries28    "config.json",29    "tokenizer_config.json",30    "preprocessor_config.json",31    "scheduler_config.json",32}33 34 35@dataclass36class DDUFEntry:37    """Object representing a file entry in a DDUF file.38 39    See [`read_dduf_file`] for how to read a DDUF file.40 41    Attributes:42        filename (str):43            The name of the file in the DDUF archive.44        offset (int):45            The offset of the file in the DDUF archive.46        length (int):47            The length of the file in the DDUF archive.48        dduf_path (str):49            The path to the DDUF archive (for internal use).50    """51 52    filename: str53    length: int54    offset: int55 56    dduf_path: Path = field(repr=False)57 58    @contextmanager59    def as_mmap(self) -> Generator[bytes, None, None]:60        """Open the file as a memory-mapped file.61 62        Useful to load safetensors directly from the file.63 64        Example:65            ```py66            >>> import safetensors.torch67            >>> with entry.as_mmap() as mm:68            ...     tensors = safetensors.torch.load(mm)69            ```70        """71        with self.dduf_path.open("rb") as f:72            with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as mm:73                yield mm[self.offset : self.offset + self.length]74 75    def read_text(self, encoding: str = "utf-8") -> str:76        """Read the file as text.77 78        Useful for '.txt' and '.json' entries.79 80        Example:81            ```py82            >>> import json83            >>> index = json.loads(entry.read_text())84            ```85        """86        with self.dduf_path.open("rb") as f:87            f.seek(self.offset)88            return f.read(self.length).decode(encoding=encoding)89 90 91def read_dduf_file(dduf_path: os.PathLike | str) -> dict[str, DDUFEntry]:92    """93    Read a DDUF file and return a dictionary of entries.94 95    Only the metadata is read, the data is not loaded in memory.96 97    Args:98        dduf_path (`str` or `os.PathLike`):99            The path to the DDUF file to read.100 101    Returns:102        `dict[str, DDUFEntry]`:103            A dictionary of [`DDUFEntry`] indexed by filename.104 105    Raises:106        - [`DDUFCorruptedFileError`]: If the DDUF file is corrupted (i.e. doesn't follow the DDUF format).107 108    Example:109        ```python110        >>> import json111        >>> import safetensors.torch112        >>> from huggingface_hub import read_dduf_file113 114        # Read DDUF metadata115        >>> dduf_entries = read_dduf_file("FLUX.1-dev.dduf")116 117        # Returns a mapping filename <> DDUFEntry118        >>> dduf_entries["model_index.json"]119        DDUFEntry(filename='model_index.json', offset=66, length=587)120 121        # Load model index as JSON122        >>> json.loads(dduf_entries["model_index.json"].read_text())123        {'_class_name': 'FluxPipeline', '_diffusers_version': '0.32.0.dev0', '_name_or_path': 'black-forest-labs/FLUX.1-dev', ...124 125        # Load VAE weights using safetensors126        >>> with dduf_entries["vae/diffusion_pytorch_model.safetensors"].as_mmap() as mm:127        ...     state_dict = safetensors.torch.load(mm)128        ```129    """130    entries = {}131    dduf_path = Path(dduf_path)132    logger.info(f"Reading DDUF file {dduf_path}")133    with zipfile.ZipFile(str(dduf_path), "r") as zf:134        for info in zf.infolist():135            logger.debug(f"Reading entry {info.filename}")136            if info.compress_type != zipfile.ZIP_STORED:137                raise DDUFCorruptedFileError("Data must not be compressed in DDUF file.")138 139            try:140                _validate_dduf_entry_name(info.filename)141            except DDUFInvalidEntryNameError as e:142                raise DDUFCorruptedFileError(f"Invalid entry name in DDUF file: {info.filename}") from e143 144            offset = _get_data_offset(zf, info)145 146            entries[info.filename] = DDUFEntry(147                filename=info.filename, offset=offset, length=info.file_size, dduf_path=dduf_path148            )149 150    # Consistency checks on the DDUF file151    if "model_index.json" not in entries:152        raise DDUFCorruptedFileError("Missing required 'model_index.json' entry in DDUF file.")153    index = json.loads(entries["model_index.json"].read_text())154    _validate_dduf_structure(index, entries.keys())155 156    logger.info(f"Done reading DDUF file {dduf_path}. Found {len(entries)} entries")157    return entries158 159 160def export_entries_as_dduf(dduf_path: str | os.PathLike, entries: Iterable[tuple[str, str | Path | bytes]]) -> None:161    """Write a DDUF file from an iterable of entries.162 163    This is a lower-level helper than [`export_folder_as_dduf`] that allows more flexibility when serializing data.164    In particular, you don't need to save the data on disk before exporting it in the DDUF file.165 166    Args:167        dduf_path (`str` or `os.PathLike`):168            The path to the DDUF file to write.169        entries (`Iterable[tuple[str, Union[str, Path, bytes]]]`):170            An iterable of entries to write in the DDUF file. Each entry is a tuple with the filename and the content.171            The filename should be the path to the file in the DDUF archive.172            The content can be a string or a pathlib.Path representing a path to a file on the local disk or directly the content as bytes.173 174    Raises:175        - [`DDUFExportError`]: If anything goes wrong during the export (e.g. invalid entry name, missing 'model_index.json', etc.).176 177    Example:178        ```python179        # Export specific files from the local disk.180        >>> from huggingface_hub import export_entries_as_dduf181        >>> export_entries_as_dduf(182        ...     dduf_path="stable-diffusion-v1-4-FP16.dduf",183        ...     entries=[ # List entries to add to the DDUF file (here, only FP16 weights)184        ...         ("model_index.json", "path/to/model_index.json"),185        ...         ("vae/config.json", "path/to/vae/config.json"),186        ...         ("vae/diffusion_pytorch_model.fp16.safetensors", "path/to/vae/diffusion_pytorch_model.fp16.safetensors"),187        ...         ("text_encoder/config.json", "path/to/text_encoder/config.json"),188        ...         ("text_encoder/model.fp16.safetensors", "path/to/text_encoder/model.fp16.safetensors"),189        ...         # ... add more entries here190        ...     ]191        ... )192        ```193 194        ```python195        # Export state_dicts one by one from a loaded pipeline196        >>> from diffusers import DiffusionPipeline197        >>> from typing import Generator, Tuple198        >>> import safetensors.torch199        >>> from huggingface_hub import export_entries_as_dduf200        >>> pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")201        ... # ... do some work with the pipeline202 203        >>> def as_entries(pipe: DiffusionPipeline) -> Generator[tuple[str, bytes], None, None]:204        ...     # Build a generator that yields the entries to add to the DDUF file.205        ...     # The first element of the tuple is the filename in the DDUF archive (must use UNIX separator!). The second element is the content of the file.206        ...     # Entries will be evaluated lazily when the DDUF file is created (only 1 entry is loaded in memory at a time)207        ...     yield "vae/config.json", pipe.vae.to_json_string().encode()208        ...     yield "vae/diffusion_pytorch_model.safetensors", safetensors.torch.save(pipe.vae.state_dict())209        ...     yield "text_encoder/config.json", pipe.text_encoder.config.to_json_string().encode()210        ...     yield "text_encoder/model.safetensors", safetensors.torch.save(pipe.text_encoder.state_dict())211        ...     # ... add more entries here212 213        >>> export_entries_as_dduf(dduf_path="stable-diffusion-v1-4.dduf", entries=as_entries(pipe))214        ```215    """216    logger.info(f"Exporting DDUF file '{dduf_path}'")217    filenames = set()218    index = None219    with zipfile.ZipFile(str(dduf_path), "w", zipfile.ZIP_STORED) as archive:220        for filename, content in entries:221            if filename in filenames:222                raise DDUFExportError(f"Can't add duplicate entry: {filename}")223            filenames.add(filename)224 225            if filename == "model_index.json":226                try:227                    index = json.loads(_load_content(content).decode())228                except json.JSONDecodeError as e:229                    raise DDUFExportError("Failed to parse 'model_index.json'.") from e230 231            try:232                filename = _validate_dduf_entry_name(filename)233            except DDUFInvalidEntryNameError as e:234                raise DDUFExportError(f"Invalid entry name: {filename}") from e235            logger.debug(f"Adding entry '{filename}' to DDUF file")236            _dump_content_in_archive(archive, filename, content)237 238    # Consistency checks on the DDUF file239    if index is None:240        raise DDUFExportError("Missing required 'model_index.json' entry in DDUF file.")241    try:242        _validate_dduf_structure(index, filenames)243    except DDUFCorruptedFileError as e:244        raise DDUFExportError("Invalid DDUF file structure.") from e245 246    logger.info(f"Done writing DDUF file {dduf_path}")247 248 249def export_folder_as_dduf(dduf_path: str | os.PathLike, folder_path: str | os.PathLike) -> None:250    """251    Export a folder as a DDUF file.252 253    AUses [`export_entries_as_dduf`] under the hood.254 255    Args:256        dduf_path (`str` or `os.PathLike`):257            The path to the DDUF file to write.258        folder_path (`str` or `os.PathLike`):259            The path to the folder containing the diffusion model.260 261    Example:262        ```python263        >>> from huggingface_hub import export_folder_as_dduf264        >>> export_folder_as_dduf(dduf_path="FLUX.1-dev.dduf", folder_path="path/to/FLUX.1-dev")265        ```266    """267    folder_path = Path(folder_path)268 269    def _iterate_over_folder() -> Iterable[tuple[str, Path]]:270        for path in Path(folder_path).glob("**/*"):271            if not path.is_file():272                continue273            if path.suffix not in DDUF_ALLOWED_ENTRIES:274                logger.debug(f"Skipping file '{path}' (file type not allowed)")275                continue276            path_in_archive = path.relative_to(folder_path)277            if len(path_in_archive.parts) >= 3:278                logger.debug(f"Skipping file '{path}' (nested directories not allowed)")279                continue280            yield path_in_archive.as_posix(), path281 282    export_entries_as_dduf(dduf_path, _iterate_over_folder())283 284 285def _dump_content_in_archive(archive: zipfile.ZipFile, filename: str, content: str | os.PathLike | bytes) -> None:286    with archive.open(filename, "w", force_zip64=True) as archive_fh:287        if isinstance(content, (str, Path)):288            content_path = Path(content)289            with content_path.open("rb") as content_fh:290                shutil.copyfileobj(content_fh, archive_fh, 1024 * 1024 * 8)  # type: ignore[misc]291        elif isinstance(content, bytes):292            archive_fh.write(content)293        else:294            raise DDUFExportError(f"Invalid content type for {filename}. Must be str, Path or bytes.")295 296 297def _load_content(content: str | Path | bytes) -> bytes:298    """Load the content of an entry as bytes.299 300    Used only for small checks (not to dump content into archive).301    """302    if isinstance(content, (str, Path)):303        return Path(content).read_bytes()304    elif isinstance(content, bytes):305        return content306    else:307        raise DDUFExportError(f"Invalid content type. Must be str, Path or bytes. Got {type(content)}.")308 309 310def _validate_dduf_entry_name(entry_name: str) -> str:311    if "." + entry_name.split(".")[-1] not in DDUF_ALLOWED_ENTRIES:312        raise DDUFInvalidEntryNameError(f"File type not allowed: {entry_name}")313    if "\\" in entry_name:314        raise DDUFInvalidEntryNameError(f"Entry names must use UNIX separators ('/'). Got {entry_name}.")315    entry_name = entry_name.strip("/")316    if entry_name.count("/") > 1:317        raise DDUFInvalidEntryNameError(f"DDUF only supports 1 level of directory. Got {entry_name}.")318    return entry_name319 320 321def _validate_dduf_structure(index: Any, entry_names: Iterable[str]) -> None:322    """323    Consistency checks on the DDUF file structure.324 325    Rules:326    - The 'model_index.json' entry is required and must contain a dictionary.327    - Each folder name must correspond to an entry in 'model_index.json'.328    - Each folder must contain at least a config file ('config.json', 'tokenizer_config.json', 'preprocessor_config.json', 'scheduler_config.json').329 330    Args:331        index (Any):332            The content of the 'model_index.json' entry.333        entry_names (Iterable[str]):334            The list of entry names in the DDUF file.335 336    Raises:337        - [`DDUFCorruptedFileError`]: If the DDUF file is corrupted (i.e. doesn't follow the DDUF format).338    """339    if not isinstance(index, dict):340        raise DDUFCorruptedFileError(f"Invalid 'model_index.json' content. Must be a dictionary. Got {type(index)}.")341 342    dduf_folders = {entry.split("/")[0] for entry in entry_names if "/" in entry}343    for folder in dduf_folders:344        if folder not in index:345            raise DDUFCorruptedFileError(f"Missing required entry '{folder}' in 'model_index.json'.")346        if not any(f"{folder}/{required_entry}" in entry_names for required_entry in DDUF_FOLDER_REQUIRED_ENTRIES):347            raise DDUFCorruptedFileError(348                f"Missing required file in folder '{folder}'. Must contains at least one of {DDUF_FOLDER_REQUIRED_ENTRIES}."349            )350 351 352def _get_data_offset(zf: zipfile.ZipFile, info: zipfile.ZipInfo) -> int:353    """354    Calculate the data offset for a file in a ZIP archive.355 356    Args:357        zf (`zipfile.ZipFile`):358            The opened ZIP file. Must be opened in read mode.359        info (`zipfile.ZipInfo`):360            The file info.361 362    Returns:363        int: The offset of the file data in the ZIP archive.364    """365    if zf.fp is None:366        raise DDUFCorruptedFileError("ZipFile object must be opened in read mode.")367 368    # Step 1: Get the local file header offset369    header_offset = info.header_offset370 371    # Step 2: Read the local file header372    zf.fp.seek(header_offset)373    local_file_header = zf.fp.read(30)  # Fixed-size part of the local header374 375    if len(local_file_header) < 30:376        raise DDUFCorruptedFileError("Incomplete local file header.")377 378    # Step 3: Parse the header fields to calculate the start of file data379    # Local file header: https://en.wikipedia.org/wiki/ZIP_(file_format)#File_headers380    filename_len = int.from_bytes(local_file_header[26:28], "little")381    extra_field_len = int.from_bytes(local_file_header[28:30], "little")382 383    # Data offset is after the fixed header, filename, and extra fields384    data_offset = header_offset + 30 + filename_len + extra_field_len385 386    return data_offset387