import torch
from comfy import model_management

dtype_list_torch = [
	"FP32",
	"FP16",
	# "FP8 E4M3", # ToDo: add these
	# "FP8 E5M2",
]

dtype_list = [
	"default",
	"auto (comfy)",
	*dtype_list_torch,
	"bnb8bit",
	"bnb4bit"
]

dtype_list_short = [
	"auto",
	*dtype_list_torch,
]

def string_to_dtype(s="none", mode=None):
	s = s.lower().strip()
	if s in ["default", "as-is"]:
		return None
	elif s in ["auto", "auto (comfy)"]:
		if mode == "vae":
			return model_management.vae_device()
		elif mode == "text_encoder":
			return model_management.text_encoder_dtype()
		elif mode == "unet":
			return model_management.unet_dtype()
		else:
			raise NotImplementedError(f"Unknown dtype mode '{mode}'")
	elif s in ["none", "auto (hf)", "auto (hf/bnb)"]:
		return None
	elif s in ["fp32", "float32", "float"]:
		return torch.float32
	elif s in ["fp16", "float16", "half"]:
		return torch.float16
	elif "fp8" in s or "float8" in s:
		if "e5m2" in s:
			return torch.float8_e5m2
		elif "e4m3" in s:
			return torch.float8_e4m3fn
		else:
			raise NotImplementedError(f"Unknown 8bit dtype '{s}'")
	elif "bnb" in s:
		assert s in ["bnb8bit", "bnb4bit"], f"Unknown bnb mode '{s}'"
		return s
	elif s is None:
		return None
	else:
		raise NotImplementedError(f"Unknown dtype '{s}'")