102 lines
3.6 KiB
Python
102 lines
3.6 KiB
Python
# This is free and unencumbered software released into the public domain.
|
|
#
|
|
# Anyone is free to copy, modify, publish, use, compile, sell, or distribute
|
|
# this software, either in source code form or as a compiled binary, for any
|
|
# purpose, commercial or non-commercial, and by any means.
|
|
#
|
|
# In jurisdictions that recognize copyright laws, the author or authors of this
|
|
# software dedicate any and all copyright interest in the software to the public
|
|
# domain. We make this dedication for the benefit of the public at large and to
|
|
# the detriment of our heirs and successors. We intend this dedication to be an
|
|
# overt act of relinquishment in perpetuity of all present and future rights to
|
|
# this software under copyright law.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
|
|
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
|
|
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
|
|
import json
|
|
|
|
import diffusers
|
|
import torch
|
|
import transformers
|
|
|
|
DEVICES = ["default", "auto", "cpu"]
|
|
DEFAULT_DEVICE = "cpu"
|
|
if torch.cuda.is_available():
|
|
DEVICES.append("cuda")
|
|
for i in range(torch.cuda.device_count()):
|
|
DEVICES.append(f"cuda:{i}")
|
|
if torch.cuda.device_count() > 0:
|
|
DEFAULT_DEVICE = "cuda:0"
|
|
DTYPES = ("default", "float32", "bfloat16", "float16", "bitsandbytes_8bit", "bitsandbytes_4bit")
|
|
|
|
def get_device(device):
|
|
if device == "default":
|
|
return DEFAULT_DEVICE
|
|
else:
|
|
return device
|
|
|
|
def mkkwargs(kwargs):
|
|
"""
|
|
Return kwargs appropriate for HuggingFace models.
|
|
"""
|
|
ret = {}
|
|
if "kwargs" in kwargs and kwargs["kwargs"] != "":
|
|
ret = json.loads(kwargs["kwargs"])
|
|
for key in kwargs:
|
|
if key == "kwargs":
|
|
continue
|
|
if key not in ret and kwargs[key] is not None:
|
|
ret[key] = kwargs[key]
|
|
return ret
|
|
|
|
def apply_device(kwargs, device, dtype, enable_model_cpu_offload=False, quant="pipeline"):
|
|
"""
|
|
Apply device and dtype properties to the kwargs. Quantizes in pipeline,
|
|
transformers, or diffusers mode. Returns the device that the result should
|
|
be moved to with `to`, or `None` if not needed.
|
|
"""
|
|
device = get_device(device)
|
|
to_device = None
|
|
if not enable_model_cpu_offload:
|
|
if ":" in device:
|
|
to_device = device
|
|
else:
|
|
kwargs["device_map"] = device
|
|
kwargs["torch_dtype"] = torch.bfloat16
|
|
|
|
if dtype[0:13] == "bitsandbytes_":
|
|
if quant == "transformers" or quant == "diffusers":
|
|
qc = {}
|
|
if dtype == "bitsandbytes_4bit":
|
|
qc["load_in_4bit"] = True
|
|
else:
|
|
qc["load_in_8bit"] = True
|
|
if quant == "transformers":
|
|
qc = transformers.BitsAndBytesConfig(**qc)
|
|
else:
|
|
qc = diffusers.BitsAndBytesConfig(**qc)
|
|
kwargs["quantization_config"] = qc
|
|
|
|
else: # "pipeline"
|
|
qc = {
|
|
"quant_backend": dtype
|
|
}
|
|
if dtype == "bitsandbytes_4bit":
|
|
qc["quant_kwargs"] = {"load_in_4bit": True}
|
|
else:
|
|
qc["quant_kwargs"] = {"load_in_8bit": True}
|
|
kwargs["quantization_config"] = diffusers.PipelineQuantizationConfig(**qc)
|
|
|
|
elif dtype == "float32":
|
|
kwargs["torch_dtype"] = torch.float32
|
|
elif dtype == "float16":
|
|
kwargs["torch_dtype"] = torch.float16
|
|
|
|
return to_device
|
|
|