31 lines
834 B
Python
31 lines
834 B
Python
# (c) TenserTensor || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
|
|
class TT_VaeEncodeContext:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"context": ("TT_CONTEXT",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TT_CONTEXT", "LATENT",)
|
|
RETURN_NAMES = ("CONTEXT", "LATENT",)
|
|
FUNCTION = "vae_encode"
|
|
CATEGORY = "TenserTensor/VAE"
|
|
|
|
def vae_encode(self, context):
|
|
vae = context["vae"]
|
|
image = context["image"]
|
|
|
|
if vae is None:
|
|
raise ValueError("VAE is required for encode")
|
|
if image is None:
|
|
raise ValueError("Pixel image is required for encode")
|
|
|
|
samples = vae.encode(image)
|
|
latent_dict = {"samples": samples}
|
|
context["latent"] = latent_dict
|
|
|
|
return (context, latent_dict,)
|