diff --git a/comfy/ldm/models/autoencoder.py b/comfy/ldm/models/autoencoder.py index e6493155ef7..13bd6e16b64 100644 --- a/comfy/ldm/models/autoencoder.py +++ b/comfy/ldm/models/autoencoder.py @@ -11,7 +11,7 @@ import comfy.ops class DiagonalGaussianRegularizer(torch.nn.Module): - def __init__(self, sample: bool = True): + def __init__(self, sample: bool = False): super().__init__() self.sample = sample @@ -19,16 +19,12 @@ def get_trainable_parameters(self) -> Any: yield from () def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]: - log = dict() posterior = DiagonalGaussianDistribution(z) if self.sample: z = posterior.sample() else: z = posterior.mode() - kl_loss = posterior.kl() - kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] - log["kl_loss"] = kl_loss - return z, log + return z, None class AbstractAutoencoder(torch.nn.Module): diff --git a/comfy/model_base.py b/comfy/model_base.py index fcdfde37859..4392355ea35 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -816,7 +816,7 @@ def extra_conds(self, **kwargs): class Flux(BaseModel): def __init__(self, model_config, model_type=ModelType.FLUX, device=None, unet_model=comfy.ldm.flux.model.Flux): super().__init__(model_config, model_type, device=device, unet_model=unet_model) - self.memory_usage_factor_conds = ("kontext",) + self.memory_usage_factor_conds = ("ref_latents",) def concat_cond(self, **kwargs): try: diff --git a/comfyui_version.py b/comfyui_version.py index 26cada11a09..c98c90499a6 100644 --- a/comfyui_version.py +++ b/comfyui_version.py @@ -1,3 +1,3 @@ # This file is automatically generated by the build process when version is # updated in pyproject.toml. -__version__ = "0.3.42" +__version__ = "0.3.43" diff --git a/pyproject.toml b/pyproject.toml index 2c6894c6e07..9d0f90032d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ComfyUI" -version = "0.3.42" +version = "0.3.43" readme = "README.md" license = { file = "LICENSE" } requires-python = ">=3.9"