From 89ff52cc8c557e6a5938537963ad9997259f8161 Mon Sep 17 00:00:00 2001
From: ExponentialML <59846140+ExponentialML@users.noreply.github.com>
Date: Thu, 21 Mar 2024 03:22:30 -0700
Subject: [PATCH] Scale latents for better color consistency

---
 nodes.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index 3fff8ef..7a15d65 100644
--- a/nodes.py
+++ b/nodes.py
@@ -28,7 +28,8 @@ class DynamiCrafterProcessor:
                 "images": ("IMAGE", ),
                 "use_interpolate": ("BOOLEAN", {"default": False}),
                 "fps": ("INT", {"default": 15, "min": 1, "max": 30, "step": 1}, ),
-                "frames": ("INT", {"default": 16})
+                "frames": ("INT", {"default": 16}),
+                "scale_latents": ("BOOLEAN", {"default": True})
             },
         }
 
@@ -151,7 +152,7 @@ class DynamiCrafterProcessor:
         frames,
     ):
         model.model_options['transformer_options']['conditioning'] = {
-            "c_concat": model.model.process_latent_in(c_concat),
+            "c_concat": c_concat,
             "image_emb": image_emb,
             'image_emb_uncond': image_emb_uncond,
             "fs": fs,
@@ -189,7 +190,8 @@ class DynamiCrafterProcessor:
         images, 
         use_interpolate,
         fps: int,
-        frames: int
+        frames: int,
+        scale_latents: bool
     ):
         self.model_patcher = model
         encoded_latent = vae.encode(images[:, :, :, :3])
@@ -201,6 +203,10 @@ class DynamiCrafterProcessor:
         image_emb_uncond = image_proj_model(encoded_image_uncond)
 
         c_concat = encoded_latent
+        c_concat = model.model.process_latent_in(c_concat)
+        if scale_latents:
+            c_concat *= 1.32
+            
         fs = torch.tensor([fps], dtype=torch.long, device=model_management.intermediate_device())
         
         model.set_model_unet_function_wrapper(self._forward)
@@ -367,4 +373,4 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DynamiCrafterLoader": "Load a DynamiCrafter Checkpoint",    
     "DynamiCrafterProcessor": "Apply DynamiCrafter",
-}
\ No newline at end of file
+}