diff --git a/invokeai/app/invocations/compel.py b/invokeai/app/invocations/compel.py index f0db3e6d9e..076ce81021 100644 --- a/invokeai/app/invocations/compel.py +++ b/invokeai/app/invocations/compel.py @@ -118,7 +118,7 @@ class CompelInvocation(BaseInvocation): conditioning_name = f"{context.graph_execution_state_id}_{self.id}_conditioning" # TODO: hacky but works ;D maybe rename latents somehow? - context.services.latents.set(conditioning_name, (c, ec)) + context.services.latents.save(conditioning_name, (c, ec)) return CompelOutput( conditioning=ConditioningField( diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index ac7139d031..64993e011a 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -20,7 +20,7 @@ from ...backend.stable_diffusion.diffusers_pipeline import ConditioningData, Sta from ...backend.stable_diffusion.schedulers import SCHEDULER_MAP from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig import numpy as np -from ..services.image_storage import ImageType +from ..services.image_file_storage import ImageType from .baseinvocation import BaseInvocation, InvocationContext from .image import ImageField, ImageOutput, build_image_output from .compel import ConditioningField @@ -144,7 +144,7 @@ class NoiseInvocation(BaseInvocation): noise = get_noise(self.width, self.height, device, self.seed) name = f'{context.graph_execution_state_id}__{self.id}' - context.services.latents.set(name, noise) + context.services.latents.save(name, noise) return build_noise_output(latents_name=name, latents=noise) @@ -260,7 +260,7 @@ class TextToLatentsInvocation(BaseInvocation): torch.cuda.empty_cache() name = f'{context.graph_execution_state_id}__{self.id}' - context.services.latents.set(name, result_latents) + context.services.latents.save(name, result_latents) return build_latents_output(latents_name=name, latents=result_latents) @@ -319,7 +319,7 @@ class LatentsToLatentsInvocation(TextToLatentsInvocation): torch.cuda.empty_cache() name = f'{context.graph_execution_state_id}__{self.id}' - context.services.latents.set(name, result_latents) + context.services.latents.save(name, result_latents) return build_latents_output(latents_name=name, latents=result_latents) @@ -404,7 +404,7 @@ class ResizeLatentsInvocation(BaseInvocation): torch.cuda.empty_cache() name = f"{context.graph_execution_state_id}__{self.id}" - context.services.latents.set(name, resized_latents) + context.services.latents.save(name, resized_latents) return build_latents_output(latents_name=name, latents=resized_latents) @@ -434,7 +434,7 @@ class ScaleLatentsInvocation(BaseInvocation): torch.cuda.empty_cache() name = f"{context.graph_execution_state_id}__{self.id}" - context.services.latents.set(name, resized_latents) + context.services.latents.save(name, resized_latents) return build_latents_output(latents_name=name, latents=resized_latents) @@ -478,5 +478,5 @@ class ImageToLatentsInvocation(BaseInvocation): ) name = f"{context.graph_execution_state_id}__{self.id}" - context.services.latents.set(name, latents) + context.services.latents.save(name, latents) return build_latents_output(latents_name=name, latents=latents) diff --git a/invokeai/app/services/latent_storage.py b/invokeai/app/services/latent_storage.py index 0184692e05..271bd17c1b 100644 --- a/invokeai/app/services/latent_storage.py +++ b/invokeai/app/services/latent_storage.py @@ -16,7 +16,7 @@ class LatentsStorageBase(ABC): pass @abstractmethod - def set(self, name: str, data: torch.Tensor) -> None: + def save(self, name: str, data: torch.Tensor) -> None: pass @abstractmethod @@ -47,7 +47,7 @@ class ForwardCacheLatentsStorage(LatentsStorageBase): self.__set_cache(name, latent) return latent - def set(self, name: str, data: torch.Tensor) -> None: + def save(self, name: str, data: torch.Tensor) -> None: self.__underlying_storage.set(name, data) self.__set_cache(name, data) @@ -80,7 +80,7 @@ class DiskLatentsStorage(LatentsStorageBase): latent_path = self.get_path(name) return torch.load(latent_path) - def set(self, name: str, data: torch.Tensor) -> None: + def save(self, name: str, data: torch.Tensor) -> None: latent_path = self.get_path(name) torch.save(data, latent_path)