update flux_model_loader node to take a T5 encoder from node field instead of hardcoded list, assume all models have been downloaded

This commit is contained in:
maryhipp 2024-08-21 14:25:53 -04:00 committed by Brandon
parent cfac7c8189
commit 5fed573a29
2 changed files with 9 additions and 4 deletions

View File

@ -49,6 +49,7 @@ class UIType(str, Enum, metaclass=MetaEnum):
ControlNetModel = "ControlNetModelField" ControlNetModel = "ControlNetModelField"
IPAdapterModel = "IPAdapterModelField" IPAdapterModel = "IPAdapterModelField"
T2IAdapterModel = "T2IAdapterModelField" T2IAdapterModel = "T2IAdapterModelField"
T5EncoderModel = "T5EncoderModelField"
SpandrelImageToImageModel = "SpandrelImageToImageModelField" SpandrelImageToImageModel = "SpandrelImageToImageModelField"
# endregion # endregion

View File

@ -177,7 +177,11 @@ class FluxModelLoaderInvocation(BaseInvocation):
input=Input.Direct, 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: def invoke(self, context: InvocationContext) -> FluxModelLoaderOutput:
model_key = self.model.key model_key = self.model.key
@ -231,9 +235,9 @@ class FluxModelLoaderInvocation(BaseInvocation):
return self._install_model( return self._install_model(
context, context,
submodel, submodel,
T5_ENCODER_MAP[self.t5_encoder]["name"], self.t5_encoder.name,
T5_ENCODER_MAP[self.t5_encoder]["repo"], "",
ModelFormat(T5_ENCODER_MAP[self.t5_encoder]["format"]), ModelFormat.T5Encoder,
ModelType.T5Encoder, ModelType.T5Encoder,
BaseModelType.Any, BaseModelType.Any,
) )