Spaces:

diffusers-internal-dev
/

diffusers-to-gguf

Paused

diffusers-to-gguf / convert_diffusion_to_gguf.py

0296fb6 about 1 year ago

13.3 kB

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations

import logging
import argparse
import json
import safetensors.torch
import os
import sys
from pathlib import Path
from typing import Any, ContextManager, cast
from torch import Tensor

import numpy as np
import torch
import gguf

# TODO: add more:
SUPPORTED_ARCHS = ["flux", "sd3", "ltxv", "hyvid", "wan", "hidream", "qwen"]

logger = logging.getLogger(__name__)


class QuantConfig:
    ftype: gguf.LlamaFileType
    qtype: gguf.GGMLQuantizationType

    def __init__(self, ftype: gguf.LlamaFileType, qtype: gguf.GGMLQuantizationType):
        self.ftype = ftype
        self.qtype = qtype


qconfig_map: dict[str, QuantConfig] = {
    "F16": QuantConfig(gguf.LlamaFileType.MOSTLY_F16, gguf.GGMLQuantizationType.F16),
    "BF16": QuantConfig(gguf.LlamaFileType.MOSTLY_BF16, gguf.GGMLQuantizationType.BF16),
    "Q8_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q8_0, gguf.GGMLQuantizationType.Q8_0),
    "Q6_K": QuantConfig(gguf.LlamaFileType.MOSTLY_Q6_K, gguf.GGMLQuantizationType.Q6_K),
    "Q5_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_K_S, gguf.GGMLQuantizationType.Q5_K),
    "Q5_1": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_1, gguf.GGMLQuantizationType.Q5_1),
    "Q5_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_0, gguf.GGMLQuantizationType.Q5_0),
    "Q4_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_K_S, gguf.GGMLQuantizationType.Q4_K),
    "Q4_1": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_1, gguf.GGMLQuantizationType.Q4_1),
    "Q4_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_0, gguf.GGMLQuantizationType.Q4_0),
    "Q3_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q3_K_S, gguf.GGMLQuantizationType.Q3_K),
    # "Q2_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q2_K, gguf.GGMLQuantizationType.Q2_K), # not yet supported in python
}


# tree of lazy tensors
class LazyTorchTensor(gguf.LazyBase):
    _tensor_type = torch.Tensor
    # to keep the type-checker happy
    dtype: torch.dtype
    shape: torch.Size

    # only used when converting a torch.Tensor to a np.ndarray
    _dtype_map: dict[torch.dtype, type] = {
        torch.float16: np.float16,
        torch.float32: np.float32,
    }

    # used for safetensors slices
    # ref: https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/src/lib.rs#L1046
    # TODO: uncomment U64, U32, and U16, ref: https://github.com/pytorch/pytorch/issues/58734
    _dtype_str_map: dict[str, torch.dtype] = {
        "F64": torch.float64,
        "F32": torch.float32,
        "BF16": torch.bfloat16,
        "F16": torch.float16,
        # "U64": torch.uint64,
        "I64": torch.int64,
        # "U32": torch.uint32,
        "I32": torch.int32,
        # "U16": torch.uint16,
        "I16": torch.int16,
        "U8": torch.uint8,
        "I8": torch.int8,
        "BOOL": torch.bool,
        "F8_E4M3": torch.float8_e4m3fn,
        "F8_E5M2": torch.float8_e5m2,
    }

    def numpy(self) -> gguf.LazyNumpyTensor:
        dtype = self._dtype_map[self.dtype]
        return gguf.LazyNumpyTensor(
            meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(dtype, self.shape),
            args=(self,),
            func=(lambda s: s.numpy()),
        )

    @classmethod
    def meta_with_dtype_and_shape(cls, dtype: torch.dtype, shape: tuple[int, ...]) -> Tensor:
        return torch.empty(size=shape, dtype=dtype, device="meta")

    @classmethod
    def from_safetensors_slice(cls, st_slice: Any) -> Tensor:
        dtype = cls._dtype_str_map[st_slice.get_dtype()]
        shape: tuple[int, ...] = tuple(st_slice.get_shape())
        lazy = cls(meta=cls.meta_with_dtype_and_shape(dtype, shape), args=(st_slice,), func=lambda s: s[:])
        return cast(torch.Tensor, lazy)

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        del types  # unused

        if kwargs is None:
            kwargs = {}

        if func is torch.Tensor.numpy:
            return args[0].numpy()

        return cls._wrap_fn(func)(*args, **kwargs)


class Converter:
    path_safetensors: Path
    endianess: gguf.GGUFEndian
    outtype: QuantConfig
    outfile: Path
    gguf_writer: gguf.GGUFWriter

    def __init__(
        self,
        arch: str,
        path_safetensors: Path,
        endianess: gguf.GGUFEndian,
        outtype: QuantConfig,
        outfile: Path,
        subfolder: str = None,
        repo_id: str = None,
        is_diffusers: bool = False,
    ):
        self.path_safetensors = path_safetensors
        self.endianess = endianess
        self.outtype = outtype
        self.outfile = outfile

        self.gguf_writer = gguf.GGUFWriter(path=None, arch=arch, endianess=self.endianess)
        self.gguf_writer.add_file_type(self.outtype.ftype)
        self.gguf_writer.add_type("diffusion")  # for HF hub to detect the type correctly
        if repo_id:
            self.gguf_writer.add_string("repo_id", repo_id)
        if subfolder:
            self.gguf_writer.add_string("subfolder", subfolder)
        if is_diffusers:
            self.gguf_writer.add_bool("is_diffusers", True)

        # load tensors and process
        from safetensors import safe_open

        ctx = cast(ContextManager[Any], safe_open(path_safetensors, framework="pt", device="cpu"))
        with ctx as model_part:
            for name in model_part.keys():
                data = model_part.get_slice(name)
                data = LazyTorchTensor.from_safetensors_slice(data)
                self.process_tensor(name, data)

    def process_tensor(self, name: str, data_torch: LazyTorchTensor) -> None:
        is_1d = len(data_torch.shape) == 1
        current_dtype = data_torch.dtype
        target_dtype = gguf.GGMLQuantizationType.F32 if is_1d else self.outtype.qtype

        if data_torch.dtype not in (torch.float16, torch.float32):
            data_torch = data_torch.to(torch.float32)

        data = data_torch.numpy()

        if current_dtype != target_dtype:
            from custom_quants import quantize as custom_quantize, QuantError

            try:
                data = custom_quantize(data, target_dtype)
            except QuantError as e:
                logger.warning("%s, %s", e, "falling back to F16")
                target_dtype = gguf.GGMLQuantizationType.F16
                data = custom_quantize(data, target_dtype)

        # reverse shape to make it similar to the internal ggml dimension order
        shape = gguf.quant_shape_from_byte_shape(data.shape, target_dtype) if data.dtype == np.uint8 else data.shape
        shape_str = f"{{{', '.join(str(n) for n in reversed(shape))}}}"
        logger.info(f"{f'%-32s' % f'{name},'} {current_dtype} --> {target_dtype.name}, shape = {shape_str}")

        # add tensor to gguf
        self.gguf_writer.add_tensor(name, data, raw_dtype=target_dtype)

    def write(self) -> None:
        self.gguf_writer.write_header_to_file(path=self.outfile)
        self.gguf_writer.write_kv_data_to_file()
        self.gguf_writer.write_tensors_to_file(progress=True)
        self.gguf_writer.close()


# https://github.com/bghira/SimpleTuner/blob/cea2457ab063f6dedb9e697830ae68a96be90641/helpers/training/save_hooks.py#L64
def _merge_sharded_checkpoints(folder: Path):
    with open(folder / "diffusion_pytorch_model.safetensors.index.json", "r") as f:
        ckpt_metadata = json.load(f)
    weight_map = ckpt_metadata.get("weight_map", None)
    if weight_map is None:
        raise KeyError("'weight_map' key not found in the shard index file.")

    # Collect all unique safetensors files from weight_map
    files_to_load = set(weight_map.values())
    merged_state_dict = {}

    # Load tensors from each unique file
    for file_name in files_to_load:
        part_file_path = folder / file_name
        if not os.path.exists(part_file_path):
            raise FileNotFoundError(f"Part file {file_name} not found.")

        with safetensors.safe_open(part_file_path, framework="pt", device="cpu") as f:
            for tensor_key in f.keys():
                if tensor_key in weight_map:
                    merged_state_dict[tensor_key] = f.get_tensor(tensor_key)

    return merged_state_dict


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Convert a flux model to GGUF")
    parser.add_argument(
        "--outfile",
        type=Path,
        default=Path("model-{ftype}.gguf"),
        help="path to write to; default: 'model-{ftype}.gguf' ; note: {ftype} will be replaced by the outtype",
    )
    parser.add_argument(
        "--outtype",
        type=str,
        choices=qconfig_map.keys(),
        default="F16",
        help="output quantization scheme",
    )
    parser.add_argument(
        "--arch",
        type=str,
        choices=SUPPORTED_ARCHS,
        help="output model architecture",
    )
    parser.add_argument(
        "--bigendian",
        action="store_true",
        help="model is executed on big endian machine",
    )
    parser.add_argument(
        "model",
        type=Path,
        help="directory containing safetensors model file",
        nargs="?",
    )
    parser.add_argument("--cache_dir", type=Path, help="Directory to store the intermediate files when needed.")
    parser.add_argument(
        "--subfolder", type=Path, default=None, help="Subfolder on the HF Hub to load checkpoints from."
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="increase output verbosity",
    )

    args = parser.parse_args()
    if args.model is None:
        parser.error("the following arguments are required: model")
    if args.arch is None:
        parser.error("the following arguments are required: --arch")
    if args.arch not in SUPPORTED_ARCHS:
        parser.error(f"Unsupported architecture: {args.arch}. Supported architectures: {', '.join(SUPPORTED_ARCHS)}")
    return args


def convert(args):
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)

    if not args.model.is_dir() and not args.model.is_file():
        if not len(str(args.model).split("/")) == 2:
            logging.error(f"Model path {args.model} does not exist.")
            sys.exit(1)

    is_diffusers = False
    repo_id = None
    merged_state_dict = None
    if args.model.is_dir():
        logging.info("Supplied a directory.")
        files = list(args.model.glob("*.safetensors"))
        n = len(files)
        if n == 0:
            logging.error("No safetensors files found.")
            sys.exit(1)
        if n == 1:
            logging.info(f"Assinging {files[0]} to `args.model`")
            args.model = files[0]
        if n > 1:
            assert args.model / "diffusion_pytorch_model.safetensors.index.json" in list(args.model.glob("*.*"))
            assert args.cache_dir
            merged_state_dict = _merge_sharded_checkpoints(args.model)
            filepath = args.cache_dir / "merged_state_dict.safetensors"
            safetensors.torch.save_file(merged_state_dict, filepath)
            logging.info(f"Serialized merged state dict to {filepath}")
            args.model = Path(filepath)

    elif len(str(args.model).split("/")) == 2:
        from huggingface_hub import snapshot_download

        logging.info("Hub repo ID detected.")
        allow_patterns = f"{args.subfolder}/*.*" if args.subfolder else None
        local_dir = snapshot_download(
            repo_id=str(args.model), local_dir=args.cache_dir, allow_patterns=allow_patterns, token=args.hf_token
        )
        repo_id = str(args.model)
        local_dir = Path(local_dir)
        local_dir = local_dir / args.subfolder if args.subfolder else local_dir
        merged_state_dict = _merge_sharded_checkpoints(local_dir)
        filepath = (
            args.cache_dir / "merged_state_dict.safetensors" if args.cache_dir else "merged_state_dict.safetensors"
        )
        safetensors.torch.save_file(merged_state_dict, filepath)
        logging.info(f"Serialized merged state dict to {filepath}")
        args.model = Path(filepath)
        is_diffusers = True

    if args.model.suffix != ".safetensors":
        logging.error(f"Model path {args.model} is not a safetensors file.")
        sys.exit(1)

    if args.outfile.suffix != ".gguf":
        logging.error("Output file must have .gguf extension.")
        sys.exit(1)

    qconfig = qconfig_map[args.outtype]
    outfile = Path(str(args.outfile).format(ftype=args.outtype.upper()))

    logger.info(f"Converting model in {args.model} to {outfile} with quantization {args.outtype}")
    converter = Converter(
        arch=args.arch,
        path_safetensors=args.model,
        endianess=gguf.GGUFEndian.BIG if args.bigendian else gguf.GGUFEndian.LITTLE,
        outtype=qconfig,
        outfile=outfile,
        repo_id=repo_id,
        subfolder=str(args.subfolder) if args.subfolder else None,
        is_diffusers=is_diffusers,
    )
    converter.write()
    logger.info(
        f"Conversion complete. Output written to {outfile}, architecture: {args.arch}, quantization: {qconfig.qtype.name}"
    )
    if merged_state_dict is not None:
        os.remove(filepath)
        logging.info(f"Removed the intermediate {filepath}.")

	#!/usr/bin/env python3
	# -- coding: utf-8 --

	from __future__ import annotations

	import logging
	import argparse
	import json
	import safetensors.torch
	import os
	import sys
	from pathlib import Path
	from typing import Any, ContextManager, cast
	from torch import Tensor

	import numpy as np
	import torch
	import gguf

	# TODO: add more:
	SUPPORTED_ARCHS = ["flux", "sd3", "ltxv", "hyvid", "wan", "hidream", "qwen"]

	logger = logging.getLogger(__name__)


	class QuantConfig:
	ftype: gguf.LlamaFileType
	qtype: gguf.GGMLQuantizationType

	def __init__(self, ftype: gguf.LlamaFileType, qtype: gguf.GGMLQuantizationType):
	self.ftype = ftype
	self.qtype = qtype


	qconfig_map: dict[str, QuantConfig] = {
	"F16": QuantConfig(gguf.LlamaFileType.MOSTLY_F16, gguf.GGMLQuantizationType.F16),
	"BF16": QuantConfig(gguf.LlamaFileType.MOSTLY_BF16, gguf.GGMLQuantizationType.BF16),
	"Q8_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q8_0, gguf.GGMLQuantizationType.Q8_0),
	"Q6_K": QuantConfig(gguf.LlamaFileType.MOSTLY_Q6_K, gguf.GGMLQuantizationType.Q6_K),
	"Q5_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_K_S, gguf.GGMLQuantizationType.Q5_K),
	"Q5_1": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_1, gguf.GGMLQuantizationType.Q5_1),
	"Q5_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_0, gguf.GGMLQuantizationType.Q5_0),
	"Q4_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_K_S, gguf.GGMLQuantizationType.Q4_K),
	"Q4_1": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_1, gguf.GGMLQuantizationType.Q4_1),
	"Q4_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_0, gguf.GGMLQuantizationType.Q4_0),
	"Q3_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q3_K_S, gguf.GGMLQuantizationType.Q3_K),
	# "Q2_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q2_K, gguf.GGMLQuantizationType.Q2_K), # not yet supported in python
	}


	# tree of lazy tensors
	class LazyTorchTensor(gguf.LazyBase):
	_tensor_type = torch.Tensor
	# to keep the type-checker happy
	dtype: torch.dtype
	shape: torch.Size

	# only used when converting a torch.Tensor to a np.ndarray
	_dtype_map: dict[torch.dtype, type] = {
	torch.float16: np.float16,
	torch.float32: np.float32,
	}

	# used for safetensors slices
	# ref: https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/src/lib.rs#L1046
	# TODO: uncomment U64, U32, and U16, ref: https://github.com/pytorch/pytorch/issues/58734
	_dtype_str_map: dict[str, torch.dtype] = {
	"F64": torch.float64,
	"F32": torch.float32,
	"BF16": torch.bfloat16,
	"F16": torch.float16,
	# "U64": torch.uint64,
	"I64": torch.int64,
	# "U32": torch.uint32,
	"I32": torch.int32,
	# "U16": torch.uint16,
	"I16": torch.int16,
	"U8": torch.uint8,
	"I8": torch.int8,
	"BOOL": torch.bool,
	"F8_E4M3": torch.float8_e4m3fn,
	"F8_E5M2": torch.float8_e5m2,
	}

	def numpy(self) -> gguf.LazyNumpyTensor:
	dtype = self._dtype_map[self.dtype]
	return gguf.LazyNumpyTensor(
	meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(dtype, self.shape),
	args=(self,),
	func=(lambda s: s.numpy()),
	)

	@classmethod
	def meta_with_dtype_and_shape(cls, dtype: torch.dtype, shape: tuple[int, ...]) -> Tensor:
	return torch.empty(size=shape, dtype=dtype, device="meta")

	@classmethod
	def from_safetensors_slice(cls, st_slice: Any) -> Tensor:
	dtype = cls._dtype_str_map[st_slice.get_dtype()]
	shape: tuple[int, ...] = tuple(st_slice.get_shape())
	lazy = cls(meta=cls.meta_with_dtype_and_shape(dtype, shape), args=(st_slice,), func=lambda s: s[:])
	return cast(torch.Tensor, lazy)

	@classmethod
	def __torch_function__(cls, func, types, args=(), kwargs=None):
	del types # unused

	if kwargs is None:
	kwargs = {}

	if func is torch.Tensor.numpy:
	return args[0].numpy()

	return cls._wrap_fn(func)(args, *kwargs)


	class Converter:
	path_safetensors: Path
	endianess: gguf.GGUFEndian
	outtype: QuantConfig
	outfile: Path
	gguf_writer: gguf.GGUFWriter

	def __init__(
	self,
	arch: str,
	path_safetensors: Path,
	endianess: gguf.GGUFEndian,
	outtype: QuantConfig,
	outfile: Path,
	subfolder: str = None,
	repo_id: str = None,
	is_diffusers: bool = False,
	):
	self.path_safetensors = path_safetensors
	self.endianess = endianess
	self.outtype = outtype
	self.outfile = outfile

	self.gguf_writer = gguf.GGUFWriter(path=None, arch=arch, endianess=self.endianess)
	self.gguf_writer.add_file_type(self.outtype.ftype)
	self.gguf_writer.add_type("diffusion") # for HF hub to detect the type correctly
	if repo_id:
	self.gguf_writer.add_string("repo_id", repo_id)
	if subfolder:
	self.gguf_writer.add_string("subfolder", subfolder)
	if is_diffusers:
	self.gguf_writer.add_bool("is_diffusers", True)

	# load tensors and process
	from safetensors import safe_open

	ctx = cast(ContextManager[Any], safe_open(path_safetensors, framework="pt", device="cpu"))
	with ctx as model_part:
	for name in model_part.keys():
	data = model_part.get_slice(name)
	data = LazyTorchTensor.from_safetensors_slice(data)
	self.process_tensor(name, data)

	def process_tensor(self, name: str, data_torch: LazyTorchTensor) -> None:
	is_1d = len(data_torch.shape) == 1
	current_dtype = data_torch.dtype
	target_dtype = gguf.GGMLQuantizationType.F32 if is_1d else self.outtype.qtype

	if data_torch.dtype not in (torch.float16, torch.float32):
	data_torch = data_torch.to(torch.float32)

	data = data_torch.numpy()

	if current_dtype != target_dtype:
	from custom_quants import quantize as custom_quantize, QuantError

	try:
	data = custom_quantize(data, target_dtype)
	except QuantError as e:
	logger.warning("%s, %s", e, "falling back to F16")
	target_dtype = gguf.GGMLQuantizationType.F16
	data = custom_quantize(data, target_dtype)

	# reverse shape to make it similar to the internal ggml dimension order
	shape = gguf.quant_shape_from_byte_shape(data.shape, target_dtype) if data.dtype == np.uint8 else data.shape
	shape_str = f"{{{', '.join(str(n) for n in reversed(shape))}}}"
	logger.info(f"{f'%-32s' % f'{name},'} {current_dtype} --> {target_dtype.name}, shape = {shape_str}")

	# add tensor to gguf
	self.gguf_writer.add_tensor(name, data, raw_dtype=target_dtype)

	def write(self) -> None:
	self.gguf_writer.write_header_to_file(path=self.outfile)
	self.gguf_writer.write_kv_data_to_file()
	self.gguf_writer.write_tensors_to_file(progress=True)
	self.gguf_writer.close()


	# https://github.com/bghira/SimpleTuner/blob/cea2457ab063f6dedb9e697830ae68a96be90641/helpers/training/save_hooks.py#L64
	def _merge_sharded_checkpoints(folder: Path):
	with open(folder / "diffusion_pytorch_model.safetensors.index.json", "r") as f:
	ckpt_metadata = json.load(f)
	weight_map = ckpt_metadata.get("weight_map", None)
	if weight_map is None:
	raise KeyError("'weight_map' key not found in the shard index file.")

	# Collect all unique safetensors files from weight_map
	files_to_load = set(weight_map.values())
	merged_state_dict = {}

	# Load tensors from each unique file
	for file_name in files_to_load:
	part_file_path = folder / file_name
	if not os.path.exists(part_file_path):
	raise FileNotFoundError(f"Part file {file_name} not found.")

	with safetensors.safe_open(part_file_path, framework="pt", device="cpu") as f:
	for tensor_key in f.keys():
	if tensor_key in weight_map:
	merged_state_dict[tensor_key] = f.get_tensor(tensor_key)

	return merged_state_dict


	def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(description="Convert a flux model to GGUF")
	parser.add_argument(
	"--outfile",
	type=Path,
	default=Path("model-{ftype}.gguf"),
	help="path to write to; default: 'model-{ftype}.gguf' ; note: {ftype} will be replaced by the outtype",
	)
	parser.add_argument(
	"--outtype",
	type=str,
	choices=qconfig_map.keys(),
	default="F16",
	help="output quantization scheme",
	)
	parser.add_argument(
	"--arch",
	type=str,
	choices=SUPPORTED_ARCHS,
	help="output model architecture",
	)
	parser.add_argument(
	"--bigendian",
	action="store_true",
	help="model is executed on big endian machine",
	)
	parser.add_argument(
	"model",
	type=Path,
	help="directory containing safetensors model file",
	nargs="?",
	)
	parser.add_argument("--cache_dir", type=Path, help="Directory to store the intermediate files when needed.")
	parser.add_argument(
	"--subfolder", type=Path, default=None, help="Subfolder on the HF Hub to load checkpoints from."
	)
	parser.add_argument(
	"--verbose",
	action="store_true",
	help="increase output verbosity",
	)

	args = parser.parse_args()
	if args.model is None:
	parser.error("the following arguments are required: model")
	if args.arch is None:
	parser.error("the following arguments are required: --arch")
	if args.arch not in SUPPORTED_ARCHS:
	parser.error(f"Unsupported architecture: {args.arch}. Supported architectures: {', '.join(SUPPORTED_ARCHS)}")
	return args


	def convert(args):
	if args.verbose:
	logging.basicConfig(level=logging.DEBUG)
	else:
	logging.basicConfig(level=logging.INFO)

	if not args.model.is_dir() and not args.model.is_file():
	if not len(str(args.model).split("/")) == 2:
	logging.error(f"Model path {args.model} does not exist.")
	sys.exit(1)

	is_diffusers = False
	repo_id = None
	merged_state_dict = None
	if args.model.is_dir():
	logging.info("Supplied a directory.")
	files = list(args.model.glob("*.safetensors"))
	n = len(files)
	if n == 0:
	logging.error("No safetensors files found.")
	sys.exit(1)
	if n == 1:
	logging.info(f"Assinging {files[0]} to `args.model`")
	args.model = files[0]
	if n > 1:
	assert args.model / "diffusion_pytorch_model.safetensors.index.json" in list(args.model.glob("."))
	assert args.cache_dir
	merged_state_dict = _merge_sharded_checkpoints(args.model)
	filepath = args.cache_dir / "merged_state_dict.safetensors"
	safetensors.torch.save_file(merged_state_dict, filepath)
	logging.info(f"Serialized merged state dict to {filepath}")
	args.model = Path(filepath)

	elif len(str(args.model).split("/")) == 2:
	from huggingface_hub import snapshot_download

	logging.info("Hub repo ID detected.")
	allow_patterns = f"{args.subfolder}/." if args.subfolder else None
	local_dir = snapshot_download(
	repo_id=str(args.model), local_dir=args.cache_dir, allow_patterns=allow_patterns, token=args.hf_token
	)
	repo_id = str(args.model)
	local_dir = Path(local_dir)
	local_dir = local_dir / args.subfolder if args.subfolder else local_dir
	merged_state_dict = _merge_sharded_checkpoints(local_dir)
	filepath = (
	args.cache_dir / "merged_state_dict.safetensors" if args.cache_dir else "merged_state_dict.safetensors"
	)
	safetensors.torch.save_file(merged_state_dict, filepath)
	logging.info(f"Serialized merged state dict to {filepath}")
	args.model = Path(filepath)
	is_diffusers = True

	if args.model.suffix != ".safetensors":
	logging.error(f"Model path {args.model} is not a safetensors file.")
	sys.exit(1)

	if args.outfile.suffix != ".gguf":
	logging.error("Output file must have .gguf extension.")
	sys.exit(1)

	qconfig = qconfig_map[args.outtype]
	outfile = Path(str(args.outfile).format(ftype=args.outtype.upper()))

	logger.info(f"Converting model in {args.model} to {outfile} with quantization {args.outtype}")
	converter = Converter(
	arch=args.arch,
	path_safetensors=args.model,
	endianess=gguf.GGUFEndian.BIG if args.bigendian else gguf.GGUFEndian.LITTLE,
	outtype=qconfig,
	outfile=outfile,
	repo_id=repo_id,
	subfolder=str(args.subfolder) if args.subfolder else None,
	is_diffusers=is_diffusers,
	)
	converter.write()
	logger.info(
	f"Conversion complete. Output written to {outfile}, architecture: {args.arch}, quantization: {qconfig.qtype.name}"
	)
	if merged_state_dict is not None:
	os.remove(filepath)
	logging.info(f"Removed the intermediate {filepath}.")