| import torch
|
| from .helper import precision_tool
|
|
|
|
|
| class set_precision:
|
| def __init__(self):
|
| pass
|
| @classmethod
|
| def INPUT_TYPES(cls):
|
| return {
|
| "required": {
|
| "latent_image": ("LATENT", ),
|
| "precision": (["16", "32", "64"], ),
|
| "set_default": ("BOOLEAN", {"default": False})
|
| },
|
| }
|
|
|
| RETURN_TYPES = ("LATENT",)
|
| RETURN_NAMES = ("passthrough",)
|
| FUNCTION = "main"
|
| CATEGORY = "RES4LYF/precision"
|
|
|
|
|
| def main(self,
|
| precision = "32",
|
| latent_image = None,
|
| set_default = False
|
| ):
|
|
|
| match precision:
|
| case "16":
|
| if set_default is True:
|
| torch.set_default_dtype(torch.float16)
|
| x = latent_image["samples"].to(torch.float16)
|
| case "32":
|
| if set_default is True:
|
| torch.set_default_dtype(torch.float32)
|
| x = latent_image["samples"].to(torch.float32)
|
| case "64":
|
| if set_default is True:
|
| torch.set_default_dtype(torch.float64)
|
| x = latent_image["samples"].to(torch.float64)
|
| return ({"samples": x}, )
|
|
|
|
|
|
|
| class set_precision_universal:
|
| def __init__(self):
|
| pass
|
| @classmethod
|
| def INPUT_TYPES(cls):
|
| return {
|
| "required": {
|
| "precision": (["bf16", "fp16", "fp32", "fp64", "passthrough"], {"default": "fp32"}),
|
| "set_default": ("BOOLEAN", {"default": False})
|
| },
|
| "optional": {
|
| "cond_pos": ("CONDITIONING",),
|
| "cond_neg": ("CONDITIONING",),
|
| "sigmas": ("SIGMAS", ),
|
| "latent_image": ("LATENT", ),
|
| },
|
| }
|
|
|
| RETURN_TYPES = ("CONDITIONING",
|
| "CONDITIONING",
|
| "SIGMAS",
|
| "LATENT",)
|
|
|
| RETURN_NAMES = ("cond_pos",
|
| "cond_neg",
|
| "sigmas",
|
| "latent_image",)
|
|
|
| FUNCTION = "main"
|
| CATEGORY = "RES4LYF/precision"
|
|
|
|
|
| def main(self,
|
| precision = "fp32",
|
| cond_pos = None,
|
| cond_neg = None,
|
| sigmas = None,
|
| latent_image = None,
|
| set_default = False
|
| ):
|
|
|
| dtype = None
|
| match precision:
|
| case "bf16":
|
| dtype = torch.bfloat16
|
| case "fp16":
|
| dtype = torch.float16
|
| case "fp32":
|
| dtype = torch.float32
|
| case "fp64":
|
| dtype = torch.float64
|
| case "passthrough":
|
| return (cond_pos, cond_neg, sigmas, latent_image, )
|
|
|
| if cond_pos is not None:
|
| cond_pos[0][0] = cond_pos[0][0].clone().to(dtype)
|
| cond_pos[0][1]["pooled_output"] = cond_pos[0][1]["pooled_output"].clone().to(dtype)
|
|
|
| if cond_neg is not None:
|
| cond_neg[0][0] = cond_neg[0][0].clone().to(dtype)
|
| cond_neg[0][1]["pooled_output"] = cond_neg[0][1]["pooled_output"].clone().to(dtype)
|
|
|
| if sigmas is not None:
|
| sigmas = sigmas.clone().to(dtype)
|
|
|
| if latent_image is not None:
|
| x = latent_image["samples"].clone().to(dtype)
|
| latent_image = {"samples": x}
|
|
|
| if set_default is True:
|
| torch.set_default_dtype(dtype)
|
|
|
| return (cond_pos, cond_neg, sigmas, latent_image, )
|
|
|
|
|
|
|
| class set_precision_advanced:
|
| def __init__(self):
|
| pass
|
| @classmethod
|
| def INPUT_TYPES(cls):
|
| return {
|
| "required": {
|
| "latent_image": ("LATENT", ),
|
| "global_precision": (["64", "32", "16"], ),
|
| "shark_precision": (["64", "32", "16"], ),
|
| },
|
| }
|
|
|
| RETURN_TYPES = ("LATENT","LATENT","LATENT","LATENT","LATENT",)
|
|
|
| RETURN_NAMES = ("passthrough",
|
| "latent_cast_to_global",
|
| "latent_16",
|
| "latent_32",
|
| "latent_64",
|
| )
|
|
|
| FUNCTION = "main"
|
| CATEGORY = "RES4LYF/precision"
|
|
|
|
|
| def main(self,
|
| global_precision = "32",
|
| shark_precision = "64",
|
| latent_image = None
|
| ):
|
|
|
| dtype_map = {
|
| "16": torch.float16,
|
| "32": torch.float32,
|
| "64": torch.float64
|
| }
|
| precision_map = {
|
| "16": 'fp16',
|
| "32": 'fp32',
|
| "64": 'fp64'
|
| }
|
|
|
| torch.set_default_dtype(dtype_map[global_precision])
|
| precision_tool.set_cast_type(precision_map[shark_precision])
|
|
|
| latent_passthrough = latent_image["samples"]
|
|
|
| latent_out16 = latent_image["samples"].to(torch.float16)
|
| latent_out32 = latent_image["samples"].to(torch.float32)
|
| latent_out64 = latent_image["samples"].to(torch.float64)
|
|
|
| target_dtype = dtype_map[global_precision]
|
| if latent_image["samples"].dtype != target_dtype:
|
| latent_image["samples"] = latent_image["samples"].to(target_dtype)
|
|
|
| latent_cast_to_global = latent_image["samples"]
|
|
|
| return ({"samples": latent_passthrough},
|
| {"samples": latent_cast_to_global},
|
| {"samples": latent_out16},
|
| {"samples": latent_out32},
|
| {"samples": latent_out64}
|
| )
|
|
|
|
|
|
|
