# huggingface/twodgirl
# License: apache-2.0

# Credits: city96.

from safetensors.torch import load_file
import gguf
import sys

filepath = sys.argv[1]
writer = gguf.GGUFWriter(filepath, arch='sd3')
target_quant = gguf.GGMLQuantizationType.Q8_0
sd_fp16 = load_file(sys.argv[2])
writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
writer.add_file_type(target_quant)
sd = {}
for key in sd_fp16.keys():
    tensor = sd_fp16[key]
    if len(tensor.shape) == 1 or len(tensor.shape) == 4:
        q = gguf.GGMLQuantizationType.F16
    else:
        q = target_quant
    sd[key] = gguf.quants.quantize(tensor.numpy(), q)
    writer.add_tensor(key, sd[key], raw_dtype=q)
writer.write_header_to_file(filepath)
writer.write_kv_data_to_file()
writer.write_tensors_to_file()
writer.close()