Spaces:
Running on Zero
Running on Zero
| import io | |
| import json | |
| import logging | |
| import os | |
| from collections.abc import Callable | |
| from typing import Any, Union | |
| import numpy as np | |
| import torch | |
| from .utils import BaseWebDataset | |
| logger = logging.getLogger(__name__) | |
| def _load_pth_bytes(raw: bytes) -> Any: | |
| """Load a torch.save-serialized object from raw bytes.""" | |
| return torch.load(io.BytesIO(raw), map_location="cpu", weights_only=False) | |
| def _coerce_latent(obj: Any) -> torch.Tensor: | |
| """Extract a latent tensor from the deserialized .pth payload. | |
| Strict mode: raises ``TypeError`` / ``KeyError`` if no tensor can be | |
| extracted. Accepts: | |
| - a bare ``torch.Tensor`` (returned as-is) | |
| - a numpy array -> converted to tensor | |
| - a dict containing one of the common keys: ``vae_latent``, | |
| ``latent``, ``latents``, ``data`` -> returns that value | |
| - a dict with exactly one tensor value -> returns it | |
| Returns: | |
| A tensor that downstream decoders expect to be 4-D ``[C, T, H, W]``; | |
| dimensionality is validated by :meth:`PytorchTensorWebDataset._decode_payload`. | |
| """ | |
| if isinstance(obj, torch.Tensor): | |
| return obj | |
| if isinstance(obj, np.ndarray): | |
| return torch.from_numpy(obj) | |
| if isinstance(obj, dict): | |
| for k in ("vae_latent", "latent", "latents", "data"): | |
| if k in obj: | |
| v = obj[k] | |
| if isinstance(v, torch.Tensor): | |
| return v | |
| if isinstance(v, np.ndarray): | |
| return torch.from_numpy(v) | |
| raise TypeError( | |
| f"_coerce_latent: payload[{k!r}] has unsupported type {type(v).__name__}" | |
| ) | |
| # Fallback: if there's exactly one tensor value, take it. | |
| tensor_vals = [v for v in obj.values() if isinstance(v, torch.Tensor)] | |
| if len(tensor_vals) == 1: | |
| return tensor_vals[0] | |
| raise KeyError( | |
| f"_coerce_latent: dict payload has no recognized latent key " | |
| f"(expected one of vae_latent/latent/latents/data); " | |
| f"keys={list(obj.keys())!r}" | |
| ) | |
| raise TypeError(f"_coerce_latent: unsupported payload type {type(obj).__name__}") | |
| class PytorchTensorWebDataset(BaseWebDataset): | |
| """Pre-encoded latent tar input, with JSONL captions as the authority. | |
| Store one <key>.latent.pth in [C,T,H,W]; yield float32 [T,C,H,W]. | |
| Optional meta.json is parsed; action payloads are ignored. The runtime has | |
| no payload-layout selector, custom transform, or disk annotation cache. | |
| target_frame_num crops a prefix and raises on short samples. | |
| """ | |
| SHARD_KEY_STRIP_PREFIX = "latent_" | |
| def __init__( | |
| self, | |
| data_path: Union[str, "os.PathLike[str]"], | |
| annotation_path: str, | |
| target_frame_num: int | None = None, | |
| shuffle: int = 1000, | |
| shuffle_initial: int = 100, | |
| shardshuffle: bool | int = True, | |
| resampled: bool = False, | |
| nodesplitter: Callable | None = None, | |
| handler: Callable | None = None, | |
| length: int | None = None, | |
| seed: int = 0, | |
| global_rank: int = 0, | |
| seed_namespace: int = 0, | |
| ): | |
| super().__init__( | |
| data_path=data_path, | |
| annotation_path=annotation_path, | |
| shuffle=shuffle, | |
| shuffle_initial=shuffle_initial, | |
| shardshuffle=shardshuffle, | |
| resampled=resampled, | |
| nodesplitter=nodesplitter, | |
| handler=handler, | |
| length=length, | |
| seed=seed, | |
| global_rank=global_rank, | |
| seed_namespace=seed_namespace, | |
| ) | |
| if isinstance(target_frame_num, bool) or ( | |
| target_frame_num is not None and not isinstance(target_frame_num, int) | |
| ): | |
| raise TypeError("target_frame_num must be a positive integer or None.") | |
| if target_frame_num is not None and target_frame_num <= 0: | |
| raise ValueError("target_frame_num must be greater than zero.") | |
| self.target_frame_num = target_frame_num | |
| logger.info( | |
| "Initialized %s with %d shard(s) " | |
| "(annotation_path=%r, caption_index_size=%d, cast_dtype=%s, " | |
| "target_frame_num=%s)", | |
| self.__class__.__name__, | |
| len(self.urls), | |
| self.annotation_path, | |
| len(self.caption_index), | |
| torch.float32, | |
| self.target_frame_num, | |
| ) | |
| # -- helpers ---------------------------------------------------------------- | |
| def _decode_payload(self, sample: dict[str, Any]) -> dict[str, Any]: | |
| """Decode one WebDataset sample into the trainer sample contract. | |
| Args: | |
| sample: Raw WebDataset sample containing a latent ``.pth`` payload, | |
| optional ``actions.json`` and optional ``meta.json``. | |
| Returns: | |
| Dict with ``prompts`` string, ``clean_latent`` float tensor | |
| ``[T, C, H, W]`` and metadata keys. On-disk latent layout must be | |
| ``[C, T, H, W]`` and is permuted before returning. | |
| Raises: | |
| ``KeyError`` / ``TypeError`` / ``ValueError`` on malformed samples. | |
| The dataset is intentionally strict so shard/index drift is visible. | |
| """ | |
| key = sample.get("__key__", "") | |
| url = sample.get("__url__", "") | |
| shard_key = self._extract_shard_key(url) | |
| latent_raw: bytes | None = None | |
| meta_obj: dict[str, Any] = {} | |
| for k, v in sample.items(): | |
| if k.startswith("__"): | |
| continue | |
| kl = k.lower() | |
| if latent_raw is None and ( | |
| kl == "latent.pth" | |
| or kl.endswith(".latent.pth") | |
| or kl == "pth" | |
| or kl.endswith(".pth") | |
| ): | |
| latent_raw = v | |
| elif kl == "meta.json" or kl.endswith(".meta.json"): | |
| meta_raw = v | |
| if isinstance(meta_raw, (bytes, bytearray)): | |
| meta_raw = meta_raw.decode("utf-8", errors="replace") | |
| meta_obj = json.loads(meta_raw) if isinstance(meta_raw, str) else (meta_raw or {}) | |
| if latent_raw is None: | |
| raise KeyError( | |
| f"{type(self).__name__}: no latent .pth payload found in " | |
| f"sample (url={url!r}, key={key!r}, fields={list(sample.keys())!r})" | |
| ) | |
| if not isinstance(latent_raw, (bytes, bytearray)): | |
| raise TypeError( | |
| f"{type(self).__name__}: latent payload has unexpected " | |
| f"type {type(latent_raw).__name__} " | |
| f"(url={url!r}, key={key!r}); expected raw bytes. " | |
| f"Did webdataset auto-decode the .pth file?" | |
| ) | |
| prompt = self._resolve_caption(shard_key, key, meta_obj) | |
| payload = _load_pth_bytes(bytes(latent_raw)) | |
| latent = _coerce_latent(payload) | |
| latent = latent.to(dtype=torch.float32) | |
| # On-disk layout is (C, T, H, W); align with the trainer layout (T, C, H, W). | |
| if latent.dim() != 4: | |
| raise ValueError( | |
| f"{type(self).__name__}: latent must be 4-D (C, T, H, W), " | |
| f"got shape {tuple(latent.shape)} (url={url!r}, key={key!r})" | |
| ) | |
| latent = latent.permute(1, 0, 2, 3).contiguous() | |
| if self.target_frame_num is not None: | |
| actual_frame_num = latent.shape[0] | |
| if actual_frame_num < self.target_frame_num: | |
| raise ValueError( | |
| f"{type(self).__name__}: latent requires at least " | |
| f"{self.target_frame_num} frame(s), got {actual_frame_num} " | |
| f"(url={url!r}, key={key!r}); short samples are not padded." | |
| ) | |
| latent = latent[: self.target_frame_num].contiguous() | |
| result: dict[str, Any] = { | |
| "prompts": prompt, | |
| "clean_latent": latent, | |
| "metadata": { | |
| "key": key, | |
| "shard_key": shard_key, | |
| }, | |
| } | |
| return result | |