ComfyUI/comfy_extras/nodes/deepfloyd/__init__.py
Benjamin Berman cff13ace64 node tweaks
2023-08-21 11:44:51 -07:00

31 lines
1.1 KiB
Python

from transformers import logging as transformers_logging
from diffusers import logging as diffusers_logging
from warnings import filterwarnings
import logging
from .deep_floyd import *
transformers_logging.set_verbosity_error()
diffusers_logging.set_verbosity_error()
logging.getLogger("xformers").addFilter(lambda r: "A matching Triton is not available" not in r.getMessage())
filterwarnings("ignore", category=FutureWarning, message="The `reduce_labels` parameter is deprecated")
filterwarnings("ignore", category=UserWarning, message="You seem to be using the pipelines sequentially on GPU")
filterwarnings("ignore", category=UserWarning, message="TypedStorage is deprecated")
NODE_CLASS_MAPPINGS = {
# DeepFloyd
"IFLoader": IFLoader,
"IFEncoder": IFEncoder,
"IFStageI": IFStageI,
"IFStageII": IFStageII,
"IFStageIII": IFStageIII,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"IFLoader": "DeepFloyd IF Loader",
"IFEncoder": "DeepFloyd IF Encoder",
"IFStageI": "DeepFloyd IF Stage I",
"IFStageII": "DeepFloyd IF Stage II",
"IFStageIII": "DeepFloyd IF Stage III",
}