wip: use step latents for reference style

__init__.py

@@ -152,6 +152,55 @@ def register_shared_norm(
 SHARE_NORM_OPTIONS = ["both", "group", "layer", "disabled"]
 SHARE_ATTN_OPTIONS = ["q+k", "q+k+v", "disabled"]
 
+class StyleAlignedSampleReferenceLatents:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
+                    "positive": ("CONDITIONING", ),
+                    "negative": ("CONDITIONING", ),
+                    "sampler": ("SAMPLER", ),
+                    "sigmas": ("SIGMAS", ),
+                    "latent_image": ("LATENT", ),
+                     }
+                }
+
+    RETURN_TYPES = ("STEP_LATENTS","LATENT")
+    RETURN_NAMES = ("ref_latents", "noised_output")
+
+    FUNCTION = "sample"
+
+    CATEGORY = "style_aligned"
+
+    def sample(self, model, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image):
+        sigmas = sigmas.flip(0)
+        if sigmas[0] == 0:
+            sigmas[0] = 0.0001
+
+        latent = latent_image
+        latent_image = latent["samples"]
+        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+
+        noise_mask = None
+        if "noise_mask" in latent:
+            noise_mask = latent["noise_mask"]
+
+        ref_latents = []
+        def callback(step: int, x0: T, x: T, steps: int):
+            ref_latents.insert(0, x[0])
+        
+        disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+        samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
+
+        out = latent.copy()
+        out["samples"] = samples
+        out_noised = out
+
+        ref_latents = torch.stack(ref_latents)
+
+        return (ref_latents, out_noised)
 
 class StyleAlignedReferenceSampler:
     @classmethod
@@ -161,7 +210,7 @@ class StyleAlignedReferenceSampler:
                 "model": ("MODEL",),
                 "share_norm": (SHARE_NORM_OPTIONS,),
                 "share_attn": (SHARE_ATTN_OPTIONS,),
-                "scale": ("FLOAT", {"default": 1, "min": 0, "max": 2.0, "step": 0.1}),
+                "scale": ("FLOAT", {"default": 1, "min": 0, "max": 2.0, "step": 0.01}),
                 "batch_size": ("INT", {"default": 2, "min": 1, "max": 8, "step": 1}),
                 "noise_seed": (
                     "INT",
@@ -181,7 +230,7 @@ class StyleAlignedReferenceSampler:
                 "negative": ("CONDITIONING",),
                 "sampler": ("SAMPLER",),
                 "sigmas": ("SIGMAS",),
-                "ref_latent": ("LATENT",),
+                "ref_latents": ("STEP_LATENTS",),
             },
         }
 
@@ -203,16 +252,16 @@ class StyleAlignedReferenceSampler:
         negative: T,
         sampler: T,
         sigmas: T,
-        ref_latent: "dict[str, T]",
+        ref_latents: T,
     ) -> "tuple[dict, dict]":
         m = model.clone()
         args = StyleAlignedArgs(share_attn)
 
         # Concat batch with style latent
-        style_latent_tensor = ref_latent["samples"]
+        style_latent_tensor = ref_latents[0].unsqueeze(0)
         height, width = style_latent_tensor.shape[-2:]
         latent_t = torch.zeros(
-            [batch_size, 4, height, width], device=ref_latent["samples"].device
+            [batch_size, 4, height, width], device=ref_latents.device
         )
         latent = {"samples": latent_t}
         noise = comfy.sample.prepare_noise(latent_t, noise_seed)
@@ -222,7 +271,13 @@ class StyleAlignedReferenceSampler:
         noise = torch.cat((ref_noise, noise), dim=0)
 
         x0_output = {}
-        callback = latent_preview.prepare_callback(m, sigmas.shape[-1] - 1, x0_output)
+        preview_callback = latent_preview.prepare_callback(m, sigmas.shape[-1] - 1, x0_output)
+
+        def callback(step: int, x0: T, x: T, steps: int):
+            preview_callback(step, x0, x, steps)
+            if (step < steps):
+                x[0] = ref_latents[step]
+                x0[0] = ref_latents[step]
 
         # Register shared norms
         share_group_norm = share_norm in ["group", "both"]
@@ -295,11 +350,13 @@ class StyleAlignedBatchAlign:
 
 NODE_CLASS_MAPPINGS = {
     "StyleAlignedReferenceSampler": StyleAlignedReferenceSampler,
+    "StyleAlignedSampleReferenceLatents": StyleAlignedSampleReferenceLatents,
     "StyleAlignedBatchAlign": StyleAlignedBatchAlign,
 }
 
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "StyleAlignedReferenceSampler": "StyleAligned Reference Sampler",
+    "StyleAlignedSampleReferenceLatents": "StyleAligned Sample Reference Latents",
     "StyleAlignedBatchAlign": "StyleAligned Batch Align",
 }