From 2ff4726347f5ed18adda8ca31df5beeba4bdec79 Mon Sep 17 00:00:00 2001
From: Shmuel Ronen <80190186+ShmuelRonen@users.noreply.github.com>
Date: Mon, 9 Dec 2024 21:22:43 +0200
Subject: [PATCH] main code

---
 nodes.py | 146 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 146 insertions(+)
 create mode 100644 nodes.py

diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..f2856fc
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,146 @@
+import torch
+from torch import nn
+import comfy.sample
+import comfy.model_management
+import comfy.utils
+import gc
+
+class MotionGuidedSampler(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.prev_styled = None
+        
+    def extract_motion_vector(self, current_latent, prev_latent):
+        return current_latent - prev_latent
+        
+    def apply_motion_vector(self, content, motion_vector, strength=0.8):
+        return content + (motion_vector * strength)
+
+    def process_frames(self, latents, model, positive, negative, seed, steps, cfg, sampler_name, scheduler, denoise):
+        batch_size = latents.shape[0]
+        processed_frames = []
+        
+        pbar = comfy.utils.ProgressBar(batch_size)
+        print(f"Processing {batch_size} frames with motion guidance")
+        
+        first_frame = comfy.sample.sample(
+            model,
+            noise=comfy.sample.prepare_noise(latents[0:1], seed, None),
+            steps=steps,
+            cfg=cfg,
+            sampler_name=sampler_name,
+            scheduler=scheduler,
+            positive=positive,
+            negative=negative,
+            latent_image=latents[0:1],
+            denoise=denoise
+        )
+        
+        processed_frames.append(first_frame)
+        self.prev_styled = first_frame
+        prev_original = latents[0:1]
+        
+        for i in range(1, batch_size):
+            current_original = latents[i:i+1]
+            
+            motion = self.extract_motion_vector(current_original, prev_original)
+            
+            motion_guided = self.apply_motion_vector(self.prev_styled, motion)
+            
+            current_frame = comfy.sample.sample(
+                model,
+                noise=comfy.sample.prepare_noise(motion_guided, seed + i, None),
+                steps=steps,
+                cfg=cfg,
+                sampler_name=sampler_name,
+                scheduler=scheduler,
+                positive=positive,
+                negative=negative,
+                latent_image=motion_guided,
+                denoise=min(0.6, denoise)  
+            )
+            
+            processed_frames.append(current_frame)
+            self.prev_styled = current_frame
+            prev_original = current_original
+            
+            pbar.update(1)
+            
+            if i % 5 == 0:
+                torch.cuda.empty_cache()
+        
+        return torch.cat(processed_frames, dim=0)
+
+    def forward(self, model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latents, denoise):
+        try:
+            samples = self.process_frames(
+                latents,
+                model,
+                positive,
+                negative,
+                noise_seed,
+                steps,
+                cfg,
+                sampler_name,
+                scheduler,
+                denoise
+            )
+            return samples
+            
+        except Exception as e:
+            print(f"Error in motion-guided sampling: {str(e)}")
+            raise e
+        finally:
+            torch.cuda.empty_cache()
+
+class MotionGuidedSamplerNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "video_latents": ("LATENT",),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
+                "sampler_name": (["euler", "euler_ancestral", "heun", "dpm_2", 
+                                "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive",
+                                "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_2m"],),
+                "scheduler": (["simple", "karras", "exponential", "normal", "ddim_uniform"],),
+                "denoise": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "motion_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
+            }
+        }
+    
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "sample"
+    CATEGORY = "sampling"
+
+    def sample(self, model, video_latents, positive, negative, seed, steps, 
+               cfg, sampler_name, scheduler, denoise, motion_strength):
+        
+        sampler = MotionGuidedSampler()
+        samples = sampler(
+            model=model,
+            noise_seed=seed,
+            steps=steps,
+            cfg=cfg,
+            sampler_name=sampler_name,
+            scheduler=scheduler,
+            positive=positive,
+            negative=negative,
+            latents=video_latents['samples'],
+            denoise=denoise
+        )
+        
+        return ({"samples": samples},)
+
+NODE_CLASS_MAPPINGS = {
+    "CohernetVideoSampler": MotionGuidedSamplerNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "CohernetVideoSampler": "Cohernet Video Sampler"
+}
\ No newline at end of file