feta for rf inversion

kijai · Dec 22, 2024 · 1614fb2 · 1614fb2
1 parent 6a10cf7
commit 1614fb2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/nodes.py b/nodes.py
@@ -139,8 +139,8 @@ def INPUT_TYPES(s):
         return {
             "required": {
                 "weight": ("FLOAT", {"default": 2.0, "min": 0, "max": 100, "step": 0.01, "tooltip": "The feta Weight of the Enhance-A-Video"}),
-                "single_blocks": ("BOOLEAN", {"default": False, "tooltip": "Enable Enhance-A-Video for single blocks"}),
-                "double_blocks": ("BOOLEAN", {"default": False, "tooltip": "Enable Enhance-A-Video for double blocks"}),
+                "single_blocks": ("BOOLEAN", {"default": True, "tooltip": "Enable Enhance-A-Video for single blocks"}),
+                "double_blocks": ("BOOLEAN", {"default": True, "tooltip": "Enable Enhance-A-Video for double blocks"}),
                 "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percentage of the steps to apply Enhance-A-Video"}),
                 "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percentage of the steps to apply Enhance-A-Video"}),
             },

diff --git a/nodes_rf_inversion.py b/nodes_rf_inversion.py
@@ -370,6 +370,14 @@ def process(self, model, hyvid_embeds, flow_shift, steps, embedded_guidance_scal
         from .latent_preview import prepare_callback
         callback = prepare_callback(transformer, steps)
 
+        if feta_args is not None:
+            set_enhance_weight(feta_args["weight"])
+            feta_start_percent = feta_args["start_percent"]
+            feta_end_percent = feta_args["end_percent"]
+            enable_enhance(feta_args["single_blocks"], feta_args["double_blocks"])
+        else:
+            disable_enhance()
+
         from comfy.utils import ProgressBar
         from tqdm import tqdm
         log.info(f"Sampling {num_frames} frames in {latents.shape[2]} latents at {width}x{height} with {len(timesteps)} inference steps")
@@ -378,6 +386,14 @@ def process(self, model, hyvid_embeds, flow_shift, steps, embedded_guidance_scal
         with tqdm(total=len(timesteps)) as progress_bar:
              for idx, (t, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
 
+                current_step_percentage = idx / len(timesteps)
+
+                if feta_args is not None:
+                    if feta_start_percent <= current_step_percentage <= feta_end_percent:
+                        enable_enhance(feta_args["single_blocks"], feta_args["double_blocks"])
+                    else:
+                        disable_enhance()
+
                 latent_model_input = latents
 
                 t_expand = t.repeat(latent_model_input.shape[0])