commit fd802682ab7761ba523b060ceeb6651f7a8687e6
Author: logtd <logoftd@gmail.com>
Date:   Thu Mar 28 18:42:47 2024 -0600

    init

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..b6f656e
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+**/__pycache__
+**/*.pyc
\ No newline at end of file
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..e69de29
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..682165d
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,14 @@
+from .nodes.apply_ref_unet_node import ApplyRefUNetNode
+from .nodes.ref_sampler_node import RefSamplerNode, RefSamplerCustomNode
+
+NODE_CLASS_MAPPINGS = {
+    "ApplyRefUNetNode": ApplyRefUNetNode,
+    "RefSamplerNode": RefSamplerNode,
+    "RefSamplerCustomNode": RefSamplerCustomNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ApplyRefUNetNode": "Apply Ref UNet",
+    "RefSamplerNode": "Ref Sampler",
+    "RefSamplerCustomNode": "Ref Sampler Custom",
+}
diff --git a/modules/attn_wrapper.py b/modules/attn_wrapper.py
new file mode 100644
index 0000000..bbaf17b
--- /dev/null
+++ b/modules/attn_wrapper.py
@@ -0,0 +1,63 @@
+import torch
+
+from .ref_controller import RefMode
+
+
+def get_attn_wrapper(cls):
+    class RefAttention(cls):
+        ref_mode = RefMode.OFF
+        ref_bank = None
+        is_cfg = True
+        style_fidelity = 0.5
+        ref_attn_weight = 0.5
+        attention_bound = 1.0
+        hidden_states = None
+
+        def forward(self, *args, **kwargs):
+            if self.ref_mode == RefMode.OFF:
+                return super().forward(*args, **kwargs)
+            
+            return self._ref_forward(*args, **kwargs)
+        
+        def _ref_forward(self, norm_hidden_states, context, **kwargs):
+            if self.ref_mode == RefMode.WRITE:
+                self.ref_bank = norm_hidden_states.detach().clone()
+                return super().forward(norm_hidden_states, context, **kwargs)
+            if self.ref_mode == RefMode.READ:
+                if self.attention_bound > self.ref_attn_weight:
+                    if len(self.ref_bank) == 2:
+                        batch_size = len(norm_hidden_states) // len(self.ref_bank)
+                        uc_bank = self.ref_bank[0].unsqueeze(0)
+                        uc_bank = torch.cat([uc_bank] * batch_size)
+                        c_bank = self.ref_bank[1].unsqueeze(0)
+                        c_bank = torch.cat([c_bank] * batch_size)
+                        bank = torch.cat([uc_bank, c_bank])
+                    else:
+                        bank = self.ref_bank
+                    bank_hidden_states = torch.cat([bank, norm_hidden_states], dim=1)
+                    attn_output_uc = (
+                        super().forward(
+                            norm_hidden_states,
+                            bank_hidden_states,
+                            **kwargs
+                        )
+                    )
+                    attn_output_c = attn_output_uc.clone()
+                    if self.is_cfg and self.style_fidelity > 0:
+                        uc_mask =  torch.Tensor(
+                                [1] * (norm_hidden_states.shape[0] // 2)
+                                + [0] * (norm_hidden_states.shape[0] // 2)
+                            ).to(norm_hidden_states.device).bool()
+                        attn_output_c[uc_mask] = (
+                            super().forward(
+                                norm_hidden_states[uc_mask],
+                                norm_hidden_states[uc_mask],
+                                **kwargs
+                            )
+                        )
+                    attn_output = self.style_fidelity * attn_output_c + (1.0 - self.style_fidelity) * attn_output_uc
+                else:
+                    attn_output = super().forward(norm_hidden_states, context, **kwargs)
+                return attn_output
+            
+    return RefAttention
diff --git a/modules/ref_controller.py b/modules/ref_controller.py
new file mode 100644
index 0000000..bce3d8e
--- /dev/null
+++ b/modules/ref_controller.py
@@ -0,0 +1,54 @@
+import enum
+
+
+class RefMode(enum.Enum):
+    OFF = 1
+    WRITE = 2
+    READ = 3
+
+
+class RefController:
+    modules = []
+    temp_modules = []
+    default_video_length = 16
+    default_full_length = 16
+
+    def add_module(self, module):
+        self.modules.append(module)
+
+    def set_mode(self, mode: RefMode):
+        for module in self.modules:
+            module.ref_mode = mode
+
+    def clear_modules(self):
+        for module in self.modules:
+            module.ref_bank = None
+        self.set_mode(RefMode.OFF)
+
+    def set_cfg(self, cfg: bool):
+        for module in self.modules:
+            module.cfg = cfg
+
+    def set_attention_bound(self, attention_bound: float):
+        for module in self.modules:
+            module.attention_bound = attention_bound
+
+    def set_style_fidelity(self, style_fidelity: float): 
+        for module in self.modules:
+            module.style_fidelity = style_fidelity
+
+    def add_temporal_transformer(self, module):
+        self.temp_modules.append(module)
+
+    def get_default_temporal_values(self):
+        if len(self.temp_modules) > 0:
+            self.default_full_length = self.temp_modules[0].full_length
+            self.default_video_length = self.temp_modules[0].video_length
+
+    def set_temporal_values_to(self, full_length, video_length):
+        for module in self.temp_modules:
+            module.full_length = full_length
+            module.video_length = video_length
+
+    def set_temporal_values_to_default(self):
+        self.set_temporal_values_to(self.default_full_length, self.default_video_length)
\ No newline at end of file
diff --git a/nodes/apply_ref_unet_node.py b/nodes/apply_ref_unet_node.py
new file mode 100644
index 0000000..5df8567
--- /dev/null
+++ b/nodes/apply_ref_unet_node.py
@@ -0,0 +1,19 @@
+
+from ..utils.module_utils import setup_ref_attn
+
+class ApplyRefUNetNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { 
+            "model": ("MODEL",),
+            "attn": (["FULL"],),
+        }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply"
+
+    CATEGORY = "flow"
+
+    def apply(self, model, attn):
+        ref_controller = setup_ref_attn(model, attn)
+        model.model_options['ref_controller'] = ref_controller
+        return (model, )
\ No newline at end of file
diff --git a/nodes/ref_sampler_node.py b/nodes/ref_sampler_node.py
new file mode 100644
index 0000000..92063f5
--- /dev/null
+++ b/nodes/ref_sampler_node.py
@@ -0,0 +1,60 @@
+from copy import deepcopy
+
+import comfy.samplers
+
+from ..utils.ref_sampling_utils import create_sampler
+from ..utils.ref_config import RefConfig
+
+
+class RefSamplerNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { 
+            "ref_latents": ("LATENT",),
+            "sampler": (comfy.samplers.SAMPLER_NAMES, ),
+            "style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
+            "attention_bound": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
+        }}
+    RETURN_TYPES = ("SAMPLER",)
+    FUNCTION = "build"
+
+    CATEGORY = "flow"
+
+    def build(self, ref_latents, sampler, style_fidelity, attention_bound):
+
+        ref_config = RefConfig(style_fidelity, attention_bound)
+
+        sampler = comfy.samplers.ksampler(sampler)
+        sampler_fn = create_sampler(sampler.sampler_function, ref_latents, ref_config)
+        
+        sampler.sampler_function = sampler_fn
+
+        return (sampler, )
+    
+
+
+
+class RefSamplerCustomNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { 
+            "ref_latents": ("LATENT",),
+            "sampler": ("SAMPLER",),
+            "style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
+            "attention_bound": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
+        }}
+    RETURN_TYPES = ("SAMPLER",)
+    FUNCTION = "build"
+
+    CATEGORY = "flow"
+
+    def build(self, ref_latents, sampler, style_fidelity, attention_bound):
+        sampler = deepcopy(sampler)
+
+        ref_config = RefConfig(style_fidelity, attention_bound)
+        
+        sampler_fn = create_sampler(sampler.sampler_function, ref_latents, ref_config)
+        
+        sampler.sampler_function = sampler_fn
+
+        return (sampler, )
\ No newline at end of file
diff --git a/utils/module_utils.py b/utils/module_utils.py
new file mode 100644
index 0000000..c1e8dea
--- /dev/null
+++ b/utils/module_utils.py
@@ -0,0 +1,73 @@
+import torch
+
+from ..modules.attn_wrapper import get_attn_wrapper
+from ..modules.ref_controller import RefController
+
+
+def isinstance_str(x: object, cls_name: str):
+    for _cls in x.__class__.__mro__:
+        if _cls.__name__ == cls_name:
+            return True
+    
+    return False
+
+
+def is_transformer_block(module):
+    return isinstance_str(module, 'BasicTransformerBlock')
+
+
+def is_temporal_block(module):
+    return isinstance_str(module, 'TemporalTransformerBlock')
+
+
+def is_named_module_transformer_block(named_module):
+    if is_transformer_block(named_module[1]):
+        return True
+    elif is_temporal_block(named_module[1]):
+        return True
+    return False
+
+
+def get_norm1_shape(block_module):
+    if is_transformer_block(block_module):
+        return block_module.norm1.normalized_shape[0]
+    elif is_temporal_block(block_module):
+        return block_module.norms[0].normalized_shape[0]
+    raise ValueError('Bad block!', block_module)
+
+
+def get_block_attention(block_module):
+    if is_transformer_block(block_module):
+        return block_module.attn1
+    elif is_temporal_block(block_module):
+        return block_module.attention_blocks[0]
+
+
+def sort_block(module):
+    return -get_norm1_shape(module)
+
+
+def setup_ref_attn(model, attn_type):
+    # TODO: attn_type for full/mid_out
+    model = model.model.diffusion_model
+    block_modules = list(filter(is_named_module_transformer_block, model.named_modules()))
+    block_modules = list(map(lambda x: x[1], block_modules))
+    block_modules = sorted(block_modules, key=sort_block)
+
+    ref_controller = RefController()
+
+    for i, block_module in enumerate(block_modules):
+        if is_transformer_block(block_module):
+            attn = get_block_attention(block_module)
+            if not hasattr(attn, 'ref_mode'):
+                attn.__class__ = get_attn_wrapper(attn.__class__)
+            attn.ref_attn_weight = float(i) / float(len(block_modules))
+            ref_controller.add_module(attn)
+
+
+    temp_transformers = list(filter(lambda x: isinstance_str(x, 'TemporalTransformer3DModel'), model.named_modules()))
+
+    for _, transformer_module in temp_transformers:
+        ref_controller.add_temporal_transformer(transformer_module)
+
+    return ref_controller
\ No newline at end of file
diff --git a/utils/noise_utils.py b/utils/noise_utils.py
new file mode 100644
index 0000000..a385afc
--- /dev/null
+++ b/utils/noise_utils.py
@@ -0,0 +1,22 @@
+import torch 
+
+
+def get_alphacumprod(sigma):
+    return 1 / ((sigma * sigma) + 1)
+
+
+def add_noise(src_latent, noise, sigma):
+    alphas_cumprod = get_alphacumprod(sigma)
+
+    sqrt_alpha_prod = alphas_cumprod ** 0.5
+    sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+    while len(sqrt_alpha_prod.shape) < len(src_latent.shape):
+        sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+    sqrt_one_minus_alpha_prod = (1 - alphas_cumprod) ** 0.5
+    sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+    while len(sqrt_one_minus_alpha_prod.shape) < len(src_latent.shape):
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+    noisy_samples = sqrt_alpha_prod * src_latent + sqrt_one_minus_alpha_prod * noise
+    return noisy_samples
diff --git a/utils/ref_config.py b/utils/ref_config.py
new file mode 100644
index 0000000..a642b68
--- /dev/null
+++ b/utils/ref_config.py
@@ -0,0 +1,7 @@
+
+
+
+class RefConfig:
+    def __init__(self, style_fidelity, attention_bound):
+        self.style_fidelity = style_fidelity
+        self.attention_bound = attention_bound
\ No newline at end of file
diff --git a/utils/ref_sampling_utils.py b/utils/ref_sampling_utils.py
new file mode 100644
index 0000000..b687a56
--- /dev/null
+++ b/utils/ref_sampling_utils.py
@@ -0,0 +1,93 @@
+import torch
+
+import comfy.model_management
+
+from .ref_config import RefConfig
+from ..modules.ref_controller import RefController, RefMode
+from ..utils.noise_utils import add_noise
+
+
+def get_base_model(model):
+    base_model = model
+    while hasattr(base_model, 'inner_model'):
+        base_model = base_model.inner_model
+    return base_model
+
+
+def prepare_ref_latents(model, ref_latent):
+    base_model = get_base_model(model)
+    ref_latent = ref_latent['samples']
+    ref_latent = base_model.process_latent_in(ref_latent)
+    device = comfy.model_management.get_torch_device()
+    ref_latent = ref_latent.to(device)
+    return ref_latent
+
+
+def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config:RefConfig):
+    # extra_args: ['cond', 'uncond', 'cond_scale', 'model_options', 'seed', 'denoise_mask']
+    ref_latent = prepare_ref_latents(model, ref_latent)
+
+    # Get ref_controller
+    model_options = extra_args.get('model_options', {})
+    ref_controller: RefController = model_options.get('ref_controller', {})
+
+    ref_conds = []
+    for cond in extra_args['cond']:
+        cond = {**cond}
+        if 'control' in cond:
+            del cond['control']
+        ref_conds.append(cond)
+
+    ref_unconds = []
+    for cond in extra_args['uncond']:
+        cond = {**cond}
+        if 'control' in cond:
+            del cond['control']
+        ref_unconds.append(cond)
+    
+    # Setup ref
+    ref_controller.clear_modules()
+    is_cfg = extra_args.get('cond_scale', 7) > 1
+    ref_controller.set_cfg(is_cfg)
+    ref_controller.set_style_fidelity(ref_config.style_fidelity)
+    ref_controller.set_attention_bound(ref_config.attention_bound)
+
+    def model_sample(latents, sigma, **kwargs):
+        ref_controller.get_default_temporal_values()
+        # Do Ref Sampling
+        ref_controller.set_temporal_values_to(1, 1)
+        ref_controller.set_mode(RefMode.WRITE)
+        ref_latent_noised = add_noise(ref_latent, torch.randn_like(ref_latent), sigma[0])
+        ref_sigma = torch.tensor([sigma[0]]*len(ref_latent)).to(sigmas.device)
+
+        model(ref_latent_noised, ref_sigma, **{**kwargs, 'cond': ref_conds, 'uncond': ref_unconds})
+
+        # Do Latent Sampling
+        ref_controller.set_temporal_values_to_default()
+        ref_controller.set_mode(RefMode.READ)
+        latents = model(latents, sigma, **kwargs)
+
+        # Clear refs
+        ref_controller.clear_modules()
+        ref_controller.set_mode(RefMode.OFF)
+        return latents
+
+
+    class RefModelWrapper:
+        inner_model = model.inner_model
+        def __call__(self, *args, **kwargs):
+            return model_sample(*args, **kwargs)
+        
+    return RefModelWrapper()
+
+
+def create_sampler(sample_fn, inversion_latent, ref_config):
+    @torch.no_grad()
+    def sample(model, latents, sigmas, extra_args=None, callback=None, disable=None, **extra_options):
+        # src_latents = prepare_src_latents(model, src_latent_image)
+        model_wrapper = create_model_wrapper(model, sigmas, inversion_latent, extra_args, ref_config)
+        output = sample_fn(model_wrapper, latents, sigmas, extra_args=extra_args, callback=callback, disable=disable)
+        del model_wrapper
+        return output
+    
+    return sample
\ No newline at end of file