Rename ControlNetFlux -> XLabsControlNetFlux

This commit is contained in:
Ryan Dick 2024-10-04 14:28:23 +00:00
parent 92b1515d9d
commit d1a0e99701
4 changed files with 8 additions and 8 deletions

View File

@ -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(

View File

@ -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:

View File

@ -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,

View File

@ -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)