mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-19 01:36:40 +08:00
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
import torch
|
|
from comfy import model_management
|
|
|
|
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 ["bf16", "bfloat16"]:
|
|
return torch.bfloat16
|
|
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}'") |