| 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}'") |