Upload 5 files

README (20).md

@@ -0,0 +1,64 @@
+---
+license: other
+license_name: flux-1-dev-non-commercial-license
+license_link: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md
+base_model:
+- ashen0209/Flux-Dev2Pro
+- black-forest-labs/FLUX.1-schnell
+tags:
+- text-to-image
+- flux
+- merge
+widget:
+- text: Example 1
+  output:
+    url: images/0015.webp
+- text: Example 2
+  output:
+    url: images/0022.webp
+- text: Example 3
+  output:
+    url: images/0012.webp
+- text: Example 4
+  output:
+    url: images/0014.webp
+- text: Example 5
+  output:
+    url: images/0033.webp
+- text: Example 6
+  output:
+    url: images/0032.webp
+---
+
+# Flux D+S F8 Diffusers
+
+A + B merge, meant for gradio inference, possibly build adapters on top with [diffusers](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_flux.py).
+
+<Gallery />
+
+## Inference
+
+[Instruction](https://huggingface.co/twodgirl/flux-dev-fp8-e4m3fn-diffusers)
+
+This model needs 10 steps, otherwise the images remain blurry.
+
+## ComfyUI
+
+I convert all my checkpoints locally. There are 5+ file formats out there, they take up too much space.
+
+Download the contents of this repo, then enter the directory. A script will convert the diffusers format to something that is compatible with ComfyUI.
+
+```
+cd flux-devpro-schnell-merge-fp8-e4m3fn-diffusers
+pip install safetensors torch
+python flux_devpro_ckpt.py transformer/diffusion_pytorch_model.safetensors flux-devpro-schnell.safetensors
+mv flux-devpro-schnell.safetensors path/to/your/ComfyUI/models/unet/flux-devpro-schnell.safetensors
+```
+
+In the workflow, the LoadDiffusionModel node should take two values: flux-devpro-schnell.safetensors, fp8_e4m3fn.
+
+Use of this code requires citation and attribution to the author via a link to their Hugging Face profile in all resulting work.
+
+## Disclaimer
+
+Sharing, reusing the model weights requires a link back to authors and their Hugging Face profile. The source models were uploaded by Ashen0209 and BFL.

flux_devpro_ckpt.py

@@ -0,0 +1,16 @@
+import gc
+from map_from_diffusers import convert_diffusers_to_flux_checkpoint
+from safetensors.torch import load_file, save_file
+import sys
+import torch
+
+###
+# Code from huggingface/twodgirl
+# License: apache-2.0
+
+if __name__ == '__main__':
+    sd = convert_diffusers_to_flux_checkpoint(load_file(sys.argv[1]))
+    assert sd['time_in.in_layer.weight'].dtype == torch.float8_e4m3fn
+    print(len(sd))
+    gc.collect()
+    save_file(sd, sys.argv[2])

gitattributes (16)

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

map_from_diffusers.py

@@ -0,0 +1,348 @@
+import torch
+
+##
+# Code from huggingface/twodgirl
+# License: apache-2.0
+#
+# Reverse of the script from
+# https://github.com/huggingface/diffusers/blob/main/scripts/convert_flux_to_diffusers.py
+
+def swap_scale_shift(weight):
+    shift, scale = weight.chunk(2, dim=0)
+    new_weight = torch.cat([scale, shift], dim=0)
+
+    return new_weight
+
+def convert_diffusers_to_flux_checkpoint(
+    converted_state_dict,
+    num_layers=19,
+    num_single_layers=38,
+    inner_dim=3072,
+    mlp_ratio=4.0
+):
+    """
+    84c3df90-9df5-48c2-9fa0-1e81324e61bf
+
+    Reverses the conversion from Diffusers checkpoint to  Flux Transformer format.
+
+    This function takes a state dictionary that has been converted to the Diffusers format
+    and transforms it back to the original Flux Transformer checkpoint format. It systematically
+    maps each parameter from the Diffusers naming and structure back to the original format,
+    handling different components such as embeddings, transformer blocks, and normalization layers.
+
+    Args:
+        converted_state_dict (dict): The state dictionary in Diffusers format to be converted back.
+        num_layers (int, optional): Number of transformer layers in the original model. Default is 19.
+        num_single_layers (int, optional): Number of single transformer layers. Default is 38.
+        inner_dim (int, optional): The inner dimension size for MLP layers. Default is 3072.
+        mlp_ratio (float, optional): The ratio to compute the MLP hidden dimension. Default is 4.0.
+
+    Returns:
+        dict: The original state dictionary in Flux Transformer checkpoint format.
+    """
+    # Initialize an empty dictionary to store the original state dictionary.
+    original_state_dict = {}
+
+    # -------------------------
+    # Handle Time Text Embeddings
+    # -------------------------
+
+    # Map the timestep embedder weights and biases back to "time_in.in_layer"
+    original_state_dict["time_in.in_layer.weight"] = converted_state_dict.pop(
+        "time_text_embed.timestep_embedder.linear_1.weight"
+    )
+    original_state_dict["time_in.in_layer.bias"] = converted_state_dict.pop(
+        "time_text_embed.timestep_embedder.linear_1.bias"
+    )
+    original_state_dict["time_in.out_layer.weight"] = converted_state_dict.pop(
+        "time_text_embed.timestep_embedder.linear_2.weight"
+    )
+    original_state_dict["time_in.out_layer.bias"] = converted_state_dict.pop(
+        "time_text_embed.timestep_embedder.linear_2.bias"
+    )
+
+    # Map the text embedder weights and biases back to "vector_in.in_layer"
+    original_state_dict["vector_in.in_layer.weight"] = converted_state_dict.pop(
+        "time_text_embed.text_embedder.linear_1.weight"
+    )
+    original_state_dict["vector_in.in_layer.bias"] = converted_state_dict.pop(
+        "time_text_embed.text_embedder.linear_1.bias"
+    )
+    original_state_dict["vector_in.out_layer.weight"] = converted_state_dict.pop(
+        "time_text_embed.text_embedder.linear_2.weight"
+    )
+    original_state_dict["vector_in.out_layer.bias"] = converted_state_dict.pop(
+        "time_text_embed.text_embedder.linear_2.bias"
+    )
+
+    # -------------------------
+    # Handle Guidance Embeddings (if present)
+    # -------------------------
+
+    # Check if any keys related to guidance are present in the converted_state_dict
+    has_guidance = any("guidance_embedder" in k for k in converted_state_dict)
+    if has_guidance:
+        # Map the guidance embedder weights and biases back to "guidance_in.in_layer"
+        original_state_dict["guidance_in.in_layer.weight"] = converted_state_dict.pop(
+            "time_text_embed.guidance_embedder.linear_1.weight"
+        )
+        original_state_dict["guidance_in.in_layer.bias"] = converted_state_dict.pop(
+            "time_text_embed.guidance_embedder.linear_1.bias"
+        )
+        original_state_dict["guidance_in.out_layer.weight"] = converted_state_dict.pop(
+            "time_text_embed.guidance_embedder.linear_2.weight"
+        )
+        original_state_dict["guidance_in.out_layer.bias"] = converted_state_dict.pop(
+            "time_text_embed.guidance_embedder.linear_2.bias"
+        )
+
+    # -------------------------
+    # Handle Context and Image Embeddings
+    # -------------------------
+
+    # Map the context embedder weights and biases back to "txt_in"
+    original_state_dict["txt_in.weight"] = converted_state_dict.pop("context_embedder.weight")
+    original_state_dict["txt_in.bias"] = converted_state_dict.pop("context_embedder.bias")
+
+    # Map the image embedder weights and biases back to "img_in"
+    original_state_dict["img_in.weight"] = converted_state_dict.pop("x_embedder.weight")
+    original_state_dict["img_in.bias"] = converted_state_dict.pop("x_embedder.bias")
+
+    # -------------------------
+    # Handle Transformer Blocks
+    # -------------------------
+
+    for i in range(num_layers):
+        # Define the prefix for the current transformer block in the converted_state_dict
+        block_prefix = f"transformer_blocks.{i}."
+
+        # -------------------------
+        # Map Norm1 Layers
+        # -------------------------
+
+        # Map the norm1 linear layer weights and biases back to "double_blocks.{i}.img_mod.lin"
+        original_state_dict[f"double_blocks.{i}.img_mod.lin.weight"] = converted_state_dict.pop(
+            f"{block_prefix}norm1.linear.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.img_mod.lin.bias"] = converted_state_dict.pop(
+            f"{block_prefix}norm1.linear.bias"
+        )
+
+        # Map the norm1_context linear layer weights and biases back to "double_blocks.{i}.txt_mod.lin"
+        original_state_dict[f"double_blocks.{i}.txt_mod.lin.weight"] = converted_state_dict.pop(
+            f"{block_prefix}norm1_context.linear.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.txt_mod.lin.bias"] = converted_state_dict.pop(
+            f"{block_prefix}norm1_context.linear.bias"
+        )
+
+        # -------------------------
+        # Handle Q, K, V Projections for Image Attention
+        # -------------------------
+
+        # Retrieve and combine the Q, K, V weights for image attention
+        q_weight = converted_state_dict.pop(f"{block_prefix}attn.to_q.weight")
+        k_weight = converted_state_dict.pop(f"{block_prefix}attn.to_k.weight")
+        v_weight = converted_state_dict.pop(f"{block_prefix}attn.to_v.weight")
+        # Concatenate along the first dimension to form the combined QKV weight
+        original_state_dict[f"double_blocks.{i}.img_attn.qkv.weight"] = torch.cat([q_weight, k_weight, v_weight], dim=0)
+
+        # Retrieve and combine the Q, K, V biases for image attention
+        q_bias = converted_state_dict.pop(f"{block_prefix}attn.to_q.bias")
+        k_bias = converted_state_dict.pop(f"{block_prefix}attn.to_k.bias")
+        v_bias = converted_state_dict.pop(f"{block_prefix}attn.to_v.bias")
+        # Concatenate along the first dimension to form the combined QKV bias
+        original_state_dict[f"double_blocks.{i}.img_attn.qkv.bias"] = torch.cat([q_bias, k_bias, v_bias], dim=0)
+
+        # -------------------------
+        # Handle Q, K, V Projections for Text Attention
+        # -------------------------
+
+        # Retrieve and combine the additional Q, K, V projections for context (text) attention
+        add_q_weight = converted_state_dict.pop(f"{block_prefix}attn.add_q_proj.weight")
+        add_k_weight = converted_state_dict.pop(f"{block_prefix}attn.add_k_proj.weight")
+        add_v_weight = converted_state_dict.pop(f"{block_prefix}attn.add_v_proj.weight")
+        # Concatenate along the first dimension to form the combined QKV weight for text
+        original_state_dict[f"double_blocks.{i}.txt_attn.qkv.weight"] = torch.cat([add_q_weight, add_k_weight, add_v_weight], dim=0)
+
+        add_q_bias = converted_state_dict.pop(f"{block_prefix}attn.add_q_proj.bias")
+        add_k_bias = converted_state_dict.pop(f"{block_prefix}attn.add_k_proj.bias")
+        add_v_bias = converted_state_dict.pop(f"{block_prefix}attn.add_v_proj.bias")
+        # Concatenate along the first dimension to form the combined QKV bias for text
+        original_state_dict[f"double_blocks.{i}.txt_attn.qkv.bias"] = torch.cat([add_q_bias, add_k_bias, add_v_bias], dim=0)
+
+        # -------------------------
+        # Map Attention Norm Layers
+        # -------------------------
+
+        # Map the attention query norm weights back to "double_blocks.{i}.img_attn.norm.query_norm.scale"
+        original_state_dict[f"double_blocks.{i}.img_attn.norm.query_norm.scale"] = converted_state_dict.pop(
+            f"{block_prefix}attn.norm_q.weight"
+        )
+
+        # Map the attention key norm weights back to "double_blocks.{i}.img_attn.norm.key_norm.scale"
+        original_state_dict[f"double_blocks.{i}.img_attn.norm.key_norm.scale"] = converted_state_dict.pop(
+            f"{block_prefix}attn.norm_k.weight"
+        )
+
+        # Map the added attention query norm weights back to "double_blocks.{i}.txt_attn.norm.query_norm.scale"
+        original_state_dict[f"double_blocks.{i}.txt_attn.norm.query_norm.scale"] = converted_state_dict.pop(
+            f"{block_prefix}attn.norm_added_q.weight"
+        )
+
+        # Map the added attention key norm weights back to "double_blocks.{i}.txt_attn.norm.key_norm.scale"
+        original_state_dict[f"double_blocks.{i}.txt_attn.norm.key_norm.scale"] = converted_state_dict.pop(
+            f"{block_prefix}attn.norm_added_k.weight"
+        )
+
+        # -------------------------
+        # Handle Feed-Forward Networks (FFNs) for Image and Text
+        # -------------------------
+
+        # Map the image MLP projection layers back to "double_blocks.{i}.img_mlp"
+        original_state_dict[f"double_blocks.{i}.img_mlp.0.weight"] = converted_state_dict.pop(
+            f"{block_prefix}ff.net.0.proj.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.img_mlp.0.bias"] = converted_state_dict.pop(
+            f"{block_prefix}ff.net.0.proj.bias"
+        )
+        original_state_dict[f"double_blocks.{i}.img_mlp.2.weight"] = converted_state_dict.pop(
+            f"{block_prefix}ff.net.2.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.img_mlp.2.bias"] = converted_state_dict.pop(
+            f"{block_prefix}ff.net.2.bias"
+        )
+
+        # Map the text MLP projection layers back to "double_blocks.{i}.txt_mlp"
+        original_state_dict[f"double_blocks.{i}.txt_mlp.0.weight"] = converted_state_dict.pop(
+            f"{block_prefix}ff_context.net.0.proj.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.txt_mlp.0.bias"] = converted_state_dict.pop(
+            f"{block_prefix}ff_context.net.0.proj.bias"
+        )
+        original_state_dict[f"double_blocks.{i}.txt_mlp.2.weight"] = converted_state_dict.pop(
+            f"{block_prefix}ff_context.net.2.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.txt_mlp.2.bias"] = converted_state_dict.pop(
+            f"{block_prefix}ff_context.net.2.bias"
+        )
+
+        # -------------------------
+        # Handle Attention Output Projections
+        # -------------------------
+
+        # Map the image attention output projection weights and biases back to "double_blocks.{i}.img_attn.proj"
+        original_state_dict[f"double_blocks.{i}.img_attn.proj.weight"] = converted_state_dict.pop(
+            f"{block_prefix}attn.to_out.0.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.img_attn.proj.bias"] = converted_state_dict.pop(
+            f"{block_prefix}attn.to_out.0.bias"
+        )
+
+        # Map the text attention output projection weights and biases back to "double_blocks.{i}.txt_attn.proj"
+        original_state_dict[f"double_blocks.{i}.txt_attn.proj.weight"] = converted_state_dict.pop(
+            f"{block_prefix}attn.to_add_out.weight"
+        )
+        original_state_dict[f"double_blocks.{i}.txt_attn.proj.bias"] = converted_state_dict.pop(
+            f"{block_prefix}attn.to_add_out.bias"
+        )
+
+    # -------------------------
+    # Handle Single Transformer Blocks
+    # -------------------------
+
+    for i in range(num_single_layers):
+        # Define the prefix for the current single transformer block in the converted_state_dict
+        block_prefix = f"single_transformer_blocks.{i}."
+
+        # -------------------------
+        # Map Norm Layers
+        # -------------------------
+
+        # Map the normalization linear layer weights and biases back to "single_blocks.{i}.modulation.lin"
+        original_state_dict[f"single_blocks.{i}.modulation.lin.weight"] = converted_state_dict.pop(
+            f"{block_prefix}norm.linear.weight"
+        )
+        original_state_dict[f"single_blocks.{i}.modulation.lin.bias"] = converted_state_dict.pop(
+            f"{block_prefix}norm.linear.bias"
+        )
+
+        # -------------------------
+        # Handle Q, K, V Projections and MLP
+        # -------------------------
+
+        # Retrieve the Q, K, V weights and the MLP projection weight
+        q_weight = converted_state_dict.pop(f"{block_prefix}attn.to_q.weight")
+        k_weight = converted_state_dict.pop(f"{block_prefix}attn.to_k.weight")
+        v_weight = converted_state_dict.pop(f"{block_prefix}attn.to_v.weight")
+        proj_mlp_weight = converted_state_dict.pop(f"{block_prefix}proj_mlp.weight")
+
+        # Concatenate Q, K, V, and MLP weights to form the combined linear1.weight
+        combined_weight = torch.cat([q_weight, k_weight, v_weight, proj_mlp_weight], dim=0)
+        original_state_dict[f"single_blocks.{i}.linear1.weight"] = combined_weight
+
+        # Retrieve the Q, K, V biases and the MLP projection bias
+        q_bias = converted_state_dict.pop(f"{block_prefix}attn.to_q.bias")
+        k_bias = converted_state_dict.pop(f"{block_prefix}attn.to_k.bias")
+        v_bias = converted_state_dict.pop(f"{block_prefix}attn.to_v.bias")
+        proj_mlp_bias = converted_state_dict.pop(f"{block_prefix}proj_mlp.bias")
+
+        # Concatenate Q, K, V, and MLP biases to form the combined linear1.bias
+        combined_bias = torch.cat([q_bias, k_bias, v_bias, proj_mlp_bias], dim=0)
+        original_state_dict[f"single_blocks.{i}.linear1.bias"] = combined_bias
+
+        # -------------------------
+        # Map Attention Normalization Weights
+        # -------------------------
+
+        # Map the attention query norm weights back to "single_blocks.{i}.norm.query_norm.scale"
+        original_state_dict[f"single_blocks.{i}.norm.query_norm.scale"] = converted_state_dict.pop(
+            f"{block_prefix}attn.norm_q.weight"
+        )
+
+        # Map the attention key norm weights back to "single_blocks.{i}.norm.key_norm.scale"
+        original_state_dict[f"single_blocks.{i}.norm.key_norm.scale"] = converted_state_dict.pop(
+            f"{block_prefix}attn.norm_k.weight"
+        )
+
+        # -------------------------
+        # Handle Projection Output
+        # -------------------------
+
+        # Map the projection output weights and biases back to "single_blocks.{i}.linear2"
+        original_state_dict[f"single_blocks.{i}.linear2.weight"] = converted_state_dict.pop(
+            f"{block_prefix}proj_out.weight"
+        )
+        original_state_dict[f"single_blocks.{i}.linear2.bias"] = converted_state_dict.pop(
+            f"{block_prefix}proj_out.bias"
+        )
+
+    # -------------------------
+    # Handle Final Output Projection and Normalization
+    # -------------------------
+
+    # Map the final output projection weights and biases back to "final_layer.linear"
+    original_state_dict["final_layer.linear.weight"] = converted_state_dict.pop("proj_out.weight")
+    original_state_dict["final_layer.linear.bias"] = converted_state_dict.pop("proj_out.bias")
+
+    # Reverse the swap_scale_shift transformation for normalization weights and biases
+    original_state_dict["final_layer.adaLN_modulation.1.weight"] = swap_scale_shift(
+        converted_state_dict.pop("norm_out.linear.weight")
+    )
+    original_state_dict["final_layer.adaLN_modulation.1.bias"] = swap_scale_shift(
+        converted_state_dict.pop("norm_out.linear.bias")
+    )
+
+    # -------------------------
+    # Handle Remaining Parameters (if any)
+    # -------------------------
+
+    # It's possible that there are remaining parameters that were not mapped.
+    # Depending on your use case, you can handle them here or raise an error.
+    if len(converted_state_dict) > 0:
+        # For debugging purposes, you might want to log or print the remaining keys
+        remaining_keys = list(converted_state_dict.keys())
+        print(f"Warning: The following keys were not mapped and remain in the state dict: {remaining_keys}")
+        # Optionally, you can choose to include them or exclude them from the original_state_dict
+
+    return original_state_dict

map_streamer.py

@@ -0,0 +1,101 @@
+import json
+import struct
+import torch
+import threading
+import warnings
+
+###
+# Code from ljleb/sd-mecha/sd_mecha/streaming.py
+
+DTYPE_MAPPING = {
+    'F64': (torch.float64, 8),
+    'F32': (torch.float32, 4),
+    'F16': (torch.float16, 2),
+    'BF16': (torch.bfloat16, 2),
+    'I8': (torch.int8, 1),
+    'I64': (torch.int64, 8),
+    'I32': (torch.int32, 4),
+    'I16': (torch.int16, 2),
+    "F8_E4M3": (torch.float8_e4m3fn, 1),
+    "F8_E5M2": (torch.float8_e5m2, 1),
+}
+
+class InSafetensorsDict:
+    def __init__(self, f, buffer_size):
+        self.default_buffer_size = buffer_size
+        self.file = f
+        self.header_size, self.header = self._read_header()
+        self.buffer = bytearray()
+        self.buffer_start_offset = 8 + self.header_size
+        self.lock = threading.Lock()
+
+    def __del__(self):
+        self.close()
+
+    def __getitem__(self, key):
+        if key not in self.header or key == "__metadata__":
+            raise KeyError(key)
+        return self._load_tensor(key)
+
+    def __iter__(self):
+        return iter(self.keys())
+
+    def __len__(self):
+        return len(self.header)
+
+    def close(self):
+        self.file.close()
+        self.buffer = None
+        self.header = None
+
+    def keys(self):
+        return (
+            key
+            for key in self.header.keys()
+            if key != "__metadata__"
+        )
+
+    def values(self):
+        for key in self.keys():
+            yield self[key]
+
+    def items(self):
+        for key in self.keys():
+            yield key, self[key]
+
+    def _read_header(self):
+        header_size_bytes = self.file.read(8)
+        header_size = struct.unpack('<Q', header_size_bytes)[0]
+        header_json = self.file.read(header_size).decode('utf-8').strip()
+        header = json.loads(header_json)
+
+        # sort by memory order to reduce seek time
+        sorted_header = dict(sorted(header.items(), key=lambda item: item[1].get('data_offsets', [0])[0]))
+        return header_size, sorted_header
+
+    def _ensure_buffer(self, start_pos, length):
+        if start_pos < self.buffer_start_offset or start_pos + length > self.buffer_start_offset + len(self.buffer):
+            self.file.seek(start_pos)
+            necessary_buffer_size = max(self.default_buffer_size, length)
+            if len(self.buffer) < necessary_buffer_size:
+                self.buffer = bytearray(necessary_buffer_size)
+            else:
+                self.buffer = self.buffer[:necessary_buffer_size]
+
+            self.file.readinto(self.buffer)
+            self.buffer_start_offset = start_pos
+
+    def _load_tensor(self, tensor_name):
+        tensor_info = self.header[tensor_name]
+        offsets = tensor_info['data_offsets']
+        dtype, dtype_bytes = DTYPE_MAPPING[tensor_info['dtype']]
+        shape = tensor_info['shape']
+        total_bytes = offsets[1] - offsets[0]
+        absolute_start_pos = 8 + self.header_size + offsets[0]
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore')
+            with self.lock:
+                self._ensure_buffer(absolute_start_pos, total_bytes)
+                buffer_offset = absolute_start_pos - self.buffer_start_offset
+                return torch.frombuffer(self.buffer, count=total_bytes // dtype_bytes, offset=buffer_offset, dtype=dtype).reshape(shape)
+