Change node input from IMAGE to LATENT

README.md

@@ -16,9 +16,9 @@ Implements the very basics of [Visual Style Prompting](https://github.com/naver-
 Clone the repository into your `custom_nodes` folder, and you'll see `Apply Visual Style Prompting` node. It should be placed between your sampler and inputs like the example image.
 This has currently only been tested with 1.5 based models.
 
-- `reference_image`: The image you wish to reference,
-- `positive`: The positive prompt to use as conditioning.
-- `reference_cond`: The prompt that describes the reference.
+- `reference_latent`: VAE-encoded image you wish to reference,
+- `positive`: Positive conditioning describing output image.
+- `reference_cond`: Conditioning describing reference image.
 - `enabled`: Enables or disables the effect. Note that this node will still be hooked even after disabling unless you remove it.
 - `denoise`: Works the same way Img2Img works, but utilized with reference and / or init images (this is experimental).
 - `input_blocks`: Focuses attention on the encoder layers.

nodes.py

@@ -11,9 +11,8 @@ class ApplyVisualStyle:
         return {
             "required": {
                 "model": ("MODEL",),
-                "clip": ("CLIP", ),
-                "vae": ("VAE", ),
-                "reference_image": ("IMAGE",),
+                "clip": ("CLIP",),
+                "reference_latent": ("LATENT",),
                 "reference_cond": ("CONDITIONING",),
                 "positive": ("CONDITIONING",),
                 "negative": ("CONDITIONING",),
@@ -56,8 +55,7 @@ class ApplyVisualStyle:
         self,
         model: comfy.model_patcher.ModelPatcher,
         clip,
-        vae,
-        reference_image,
+        reference_latent,
         reference_cond,
         positive,
         negative,
@@ -68,8 +66,7 @@ class ApplyVisualStyle:
         output_blocks,
         init_image = None
     ):
-        reference_latent = vae.encode(reference_image[:,:,:,:3])
-
+        reference_samples = reference_latent["samples"]
         block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks)
 
         for n, m in model.model.diffusion_model.named_modules():
@@ -85,13 +82,13 @@ class ApplyVisualStyle:
         positive_cat = cat_cond(clip, reference_cond, positive)
         negative_cat = cat_cond(clip, negative, negative)
 
-        latents = torch.zeros_like(reference_latent)
+        latents = torch.zeros_like(reference_samples)
         latents = torch.cat([latents] * 2)
 
         if denoise < 1.0:
-            latents[::1] = reference_latent[:1]
+            latents[::1] = reference_samples[:1]
         else:
-            latents[::2] = reference_latent
+            latents[::2] = reference_samples
 
         denoise_mask = torch.ones_like(latents)[:, :1, ...] * denoise