diff --git a/invokeai/app/invocations/compel.py b/invokeai/app/invocations/compel.py index 0103e3af55..4c1a4448de 100644 --- a/invokeai/app/invocations/compel.py +++ b/invokeai/app/invocations/compel.py @@ -8,8 +8,7 @@ from .model import ClipField from ...backend.util.devices import torch_dtype from ...backend.stable_diffusion.diffusion import InvokeAIDiffuserComponent -from ...backend.model_management import BaseModelType, ModelType, SubModelType -from ...backend.model_management.lora import ModelPatcher +from ...backend.model_management import BaseModelType, ModelType, SubModelType, ModelPatcher from compel import Compel from compel.prompt_parser import ( diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index a8b9131775..942dae77f0 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -24,7 +24,7 @@ from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import \ PostprocessingSettings from ...backend.stable_diffusion.schedulers import SCHEDULER_MAP from ...backend.util.devices import choose_torch_device, torch_dtype -from ...backend.model_management.lora import ModelPatcher +from ...backend.model_management import ModelPatcher from .baseinvocation import (BaseInvocation, BaseInvocationOutput, InvocationConfig, InvocationContext) from .compel import ConditioningField diff --git a/invokeai/app/invocations/onnx.py b/invokeai/app/invocations/onnx.py index 12a928c849..5d621ab9b4 100644 --- a/invokeai/app/invocations/onnx.py +++ b/invokeai/app/invocations/onnx.py @@ -14,7 +14,7 @@ from diffusers.image_processor import VaeImageProcessor from diffusers.schedulers import SchedulerMixin as Scheduler from ..models.image import ImageCategory, ImageField, ResourceOrigin -from ...backend.model_management.lora import ONNXModelPatcher +from ...backend.model_management import ONNXModelPatcher from .baseinvocation import (BaseInvocation, BaseInvocationOutput, InvocationConfig, InvocationContext) from .compel import ConditioningField @@ -24,6 +24,7 @@ from .model import ModelInfo, UNetField, VaeField from invokeai.backend import BaseModelType, ModelType, SubModelType +from tqdm import tqdm from .model import ClipField from .latent import LatentsField, LatentsOutput, build_latents_output, get_scheduler, SAMPLER_NAME_VALUES from .compel import CompelOutput @@ -88,6 +89,8 @@ class ONNXPromptInvocation(BaseInvocation): text_encoder.create_session() + # copy from + # https://github.com/huggingface/diffusers/blob/3ebbaf7c96801271f9e6c21400033b6aa5ffcf29/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L153 text_inputs = tokenizer( self.prompt, padding="max_length", @@ -169,6 +172,8 @@ class ONNXTextToLatentsInvocation(BaseInvocation): }, } + # based on + # https://github.com/huggingface/diffusers/blob/3ebbaf7c96801271f9e6c21400033b6aa5ffcf29/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L375 def invoke(self, context: InvocationContext) -> LatentsOutput: c, _ = context.services.latents.get(self.positive_conditioning.conditioning_name) uc, _ = context.services.latents.get(self.negative_conditioning.conditioning_name) @@ -224,7 +229,6 @@ class ONNXTextToLatentsInvocation(BaseInvocation): ) timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] - from tqdm import tqdm for i in tqdm(range(len(scheduler.timesteps))): t = scheduler.timesteps[i] # expand the latents if we are doing classifier free guidance @@ -260,24 +264,6 @@ class ONNXTextToLatentsInvocation(BaseInvocation): context.services.latents.save(name, latents) return build_latents_output(latents_name=name, latents=latents) - -@staticmethod -def numpy_to_pil(images): - """ - Convert a numpy image or a batch of images to a PIL image. - """ - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - if images.shape[-1] == 1: - # special case for grayscale (single channel) images - pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] - else: - pil_images = [Image.fromarray(image) for image in images] - - return pil_images - - # Latent to image class ONNXLatentsToImageInvocation(BaseInvocation): """Generates an image from latents.""" @@ -311,9 +297,10 @@ class ONNXLatentsToImageInvocation(BaseInvocation): torch.cuda.empty_cache() with vae_info as vae: - vae.create_session() + # copied from + # https://github.com/huggingface/diffusers/blob/3ebbaf7c96801271f9e6c21400033b6aa5ffcf29/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L427 latents = 1 / 0.18215 * latents # image = self.vae_decoder(latent_sample=latents)[0] # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 @@ -323,13 +310,10 @@ class ONNXLatentsToImageInvocation(BaseInvocation): image = np.clip(image / 2 + 0.5, 0, 1) image = image.transpose((0, 2, 3, 1)) - image = VaeImageProcessor.numpy_to_pil(image)[0] vae.release_session() - - torch.cuda.empty_cache() image_dto = context.services.images.create( diff --git a/invokeai/backend/model_management/__init__.py b/invokeai/backend/model_management/__init__.py index aea7b417a1..e96b064658 100644 --- a/invokeai/backend/model_management/__init__.py +++ b/invokeai/backend/model_management/__init__.py @@ -4,3 +4,4 @@ Initialization file for invokeai.backend.model_management from .model_manager import ModelManager, ModelInfo from .model_cache import ModelCache from .models import BaseModelType, ModelType, SubModelType, ModelVariantType +from .lora import ModelPatcher, ONNXModelPatcher diff --git a/invokeai/backend/model_management/lora.py b/invokeai/backend/model_management/lora.py index 6f64141610..c092ecb384 100644 --- a/invokeai/backend/model_management/lora.py +++ b/invokeai/backend/model_management/lora.py @@ -16,6 +16,8 @@ import numpy as np from compel.embeddings_provider import BaseTextualInversionManager +# TODO: rename and split this file + class LoRALayerBase: #rank: Optional[int] #alpha: Optional[float] @@ -702,7 +704,8 @@ class ONNXModelPatcher: with cls.apply_lora(text_encoder, loras, "lora_te_"): yield - + # based on + # https://github.com/ssube/onnx-web/blob/ca2e436f0623e18b4cfe8a0363fcfcf10508acf7/api/onnx_web/convert/diffusion/lora.py#L323 @classmethod @contextmanager def apply_lora( diff --git a/invokeai/backend/model_management/models/base.py b/invokeai/backend/model_management/models/base.py index 0b22e380ba..dd8fe2ee19 100644 --- a/invokeai/backend/model_management/models/base.py +++ b/invokeai/backend/model_management/models/base.py @@ -459,6 +459,7 @@ class IAIOnnxRuntimeModel(OnnxRuntimeModel): if self.session is None: #onnx.save(self.proto, "tmp.onnx") #onnx.save_model(self.proto, "tmp.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="tmp.onnx_data", size_threshold=1024, convert_attribute=False) + # TODO: something to be able to get weight when they already moved outside of model proto (trimmed_model, external_data) = buffer_external_data_tensors(self.proto) sess = SessionOptions() self._external_data.update(**external_data)