diff --git a/invokeai/app/invocations/flux_denoise.py b/invokeai/app/invocations/flux_denoise.py index d00f1d415c..ca593994a1 100644 --- a/invokeai/app/invocations/flux_denoise.py +++ b/invokeai/app/invocations/flux_denoise.py @@ -20,7 +20,7 @@ from invokeai.app.invocations.fields import ( from invokeai.app.invocations.model import TransformerField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.flux.controlnet.controlnet_flux import ControlNetFlux +from invokeai.backend.flux.controlnet.xlabs_controlnet_flux import XLabsControlNetFlux from invokeai.backend.flux.controlnet_extension import ControlNetExtension from invokeai.backend.flux.denoise import denoise from invokeai.backend.flux.inpaint_extension import InpaintExtension @@ -328,7 +328,7 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard): controlnet_extensions: list[ControlNetExtension] = [] for controlnet in controlnets: model = exit_stack.enter_context(context.models.load(controlnet.control_model)) - assert isinstance(model, ControlNetFlux) + assert isinstance(model, XLabsControlNetFlux) image = context.images.get_pil(controlnet.image.image_name) controlnet_extensions.append( diff --git a/invokeai/backend/flux/controlnet/controlnet_flux.py b/invokeai/backend/flux/controlnet/xlabs_controlnet_flux.py similarity index 99% rename from invokeai/backend/flux/controlnet/controlnet_flux.py rename to invokeai/backend/flux/controlnet/xlabs_controlnet_flux.py index 90553df9a5..00876394ee 100644 --- a/invokeai/backend/flux/controlnet/controlnet_flux.py +++ b/invokeai/backend/flux/controlnet/xlabs_controlnet_flux.py @@ -17,7 +17,7 @@ def _zero_module(module: torch.nn.Module) -> torch.nn.Module: return module -class ControlNetFlux(nn.Module): +class XLabsControlNetFlux(nn.Module): """A ControlNet model for FLUX. The architecture is very similar to the base FLUX model, with the following differences: diff --git a/invokeai/backend/flux/controlnet_extension.py b/invokeai/backend/flux/controlnet_extension.py index 39c71a1883..3b8b37ca36 100644 --- a/invokeai/backend/flux/controlnet_extension.py +++ b/invokeai/backend/flux/controlnet_extension.py @@ -6,13 +6,13 @@ from PIL.Image import Image from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES, prepare_control_image -from invokeai.backend.flux.controlnet.controlnet_flux import ControlNetFlux +from invokeai.backend.flux.controlnet.xlabs_controlnet_flux import XLabsControlNetFlux class ControlNetExtension: def __init__( self, - model: ControlNetFlux, + model: XLabsControlNetFlux, controlnet_cond: torch.Tensor, weight: Union[float, List[float]], begin_step_percent: float, @@ -30,7 +30,7 @@ class ControlNetExtension: @classmethod def from_controlnet_image( cls, - model: ControlNetFlux, + model: XLabsControlNetFlux, controlnet_image: Image, latent_height: int, latent_width: int, diff --git a/invokeai/backend/model_manager/load/model_loaders/flux.py b/invokeai/backend/model_manager/load/model_loaders/flux.py index 01e7dd32bd..1b954bdbc3 100644 --- a/invokeai/backend/model_manager/load/model_loaders/flux.py +++ b/invokeai/backend/model_manager/load/model_loaders/flux.py @@ -10,7 +10,7 @@ from safetensors.torch import load_file from transformers import AutoConfig, AutoModelForTextEncoding, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer from invokeai.app.services.config.config_default import get_config -from invokeai.backend.flux.controlnet.controlnet_flux import ControlNetFlux +from invokeai.backend.flux.controlnet.xlabs_controlnet_flux import XLabsControlNetFlux from invokeai.backend.flux.model import Flux from invokeai.backend.flux.modules.autoencoder import AutoEncoder from invokeai.backend.flux.util import ae_params, params @@ -311,7 +311,7 @@ class FluxControlnetModel(ModelLoader): with accelerate.init_empty_weights(): # HACK(ryand): Is it safe to assume dev here? - model = ControlNetFlux(params["flux-dev"]) + model = XLabsControlNetFlux(params["flux-dev"]) sd = load_file(model_path) model.load_state_dict(sd, assign=True)