From 2c72295b1cb35d49b38610d65bffcf9abacfca69 Mon Sep 17 00:00:00 2001 From: maryhipp Date: Wed, 21 Aug 2024 14:25:53 -0400 Subject: [PATCH] update flux_model_loader node to take a T5 encoder from node field instead of hardcoded list, assume all models have been downloaded --- invokeai/app/invocations/fields.py | 1 + invokeai/app/invocations/model.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py index 1b52f27fb5..6b7d7bef63 100644 --- a/invokeai/app/invocations/fields.py +++ b/invokeai/app/invocations/fields.py @@ -49,6 +49,7 @@ class UIType(str, Enum, metaclass=MetaEnum): ControlNetModel = "ControlNetModelField" IPAdapterModel = "IPAdapterModelField" T2IAdapterModel = "T2IAdapterModelField" + T5EncoderModel = "T5EncoderModelField" SpandrelImageToImageModel = "SpandrelImageToImageModelField" # endregion diff --git a/invokeai/app/invocations/model.py b/invokeai/app/invocations/model.py index 984266239c..300e38b231 100644 --- a/invokeai/app/invocations/model.py +++ b/invokeai/app/invocations/model.py @@ -177,7 +177,11 @@ class FluxModelLoaderInvocation(BaseInvocation): input=Input.Direct, ) - t5_encoder: T5_ENCODER_OPTIONS = InputField(description="The T5 Encoder model to use.") + t5_encoder: ModelIdentifierField = InputField( + description=FieldDescriptions.t5Encoder, + ui_type=UIType.T5EncoderModel, + input=Input.Direct, + ) def invoke(self, context: InvocationContext) -> FluxModelLoaderOutput: model_key = self.model.key @@ -231,9 +235,9 @@ class FluxModelLoaderInvocation(BaseInvocation): return self._install_model( context, submodel, - T5_ENCODER_MAP[self.t5_encoder]["name"], - T5_ENCODER_MAP[self.t5_encoder]["repo"], - ModelFormat(T5_ENCODER_MAP[self.t5_encoder]["format"]), + self.t5_encoder.name, + "", + ModelFormat.T5Encoder, ModelType.T5Encoder, BaseModelType.Any, )