Merge pull request #5 from ExponentialML/feat/transformer-controlnet

Add support for transformer options and ControlNet

README.md

@@ -56,15 +56,32 @@ If you are downloading the CLIP and VAE models separately, place them under thei
 - **empty_latent**: An empty latent with the same size and frames as the processed ones.
   
 - **latent_img**: If you're doing Img2Img based workflows, this is the necessary one to use.
-  
+
+# ControlNet Support
+
+You can now use DynamiCrafter by applying ControlNet to the Spatial (image) portion to guide video generations in various ways.
+The ControlNets are based on 2.1, so you must download them at the link below (Thanks @thibaud !) .
+
+**ControlNet 2.1**: https://huggingface.co/thibaud/controlnet-sd21
+
+After you download them, you can use them as you would with any other workflow.
+
+# Tips
+
 > [!TIP]
 > You don't have to use the latent outputs. As long as you use the same frame length (as your batch size) and same height and with as your image inputs, you can use your own latents.
 > This means that you can experiment with inpainting and so on.
 
+> [!TIP]
+> You can choose which frame you use as init by using VAE Encode Inpaint or Set Latent Noise Mask. You set the beginning batch mask to full black, while the rest are at full white. This also means you can do interpolation with regular models.
+> As these workflows are more advanced, examples will arrive at a future date.
+
 # TODO
 - [x] Add various workflows.
-- [ ] Add support for Spatial Transformer options.
-- [ ] Ensure attention optimizations are working properly.
+- [ ] Add advanced workflows.
+- [x] Add support for Spatial Transformer options.
+- [x] Add ControlNet support.
+- [x] Ensure attention optimizations are working properly.
 - [ ] Add autoregressive nodes (this may be a separate repository)
 - [x] Add examples. (For more, [check here](https://github.com/Doubiiu/DynamiCrafter?tab=readme-ov-file#11-showcases-576x1024)).
 

lvdm/modules/networks/openaimodel3d.py

@@ -15,6 +15,8 @@ from ...basics import (
 )
 from ...modules.attention import SpatialTransformer, TemporalTransformer
 import comfy.ops
+import logging
+
 ops = comfy.ops.disable_weight_init
 
 class TimestepBlock(nn.Module):
@@ -28,20 +30,20 @@ class TimestepBlock(nn.Module):
         """
 
 #This is needed because accelerate makes a copy of transformer_options which breaks "transformer_index"
-def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None, time_context=None, num_video_frames=None, image_only_indicator=None):
+def forward_timestep_embed(ts, x, emb, context=None, batch_size=None, transformer_options={}):
     for layer in ts:
         if isinstance(layer, TimestepBlock):
-            x = layer(x, emb)
-        elif isinstance(layer, SpatialVideoTransformer):
-            x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options)
+            x = layer(x, emb, batch_size=batch_size)
+        elif isinstance(layer, SpatialTransformer):
+            x = layer(x, context)
             if "transformer_index" in transformer_options:
                 transformer_options["transformer_index"] += 1
         elif isinstance(layer, TemporalTransformer):
-            x = layer(x, context, transformer_options)
+            x = rearrange(x, '(b f) c h w -> b c f h w', b=batch_size)
+            x = layer(x, context)
             if "transformer_index" in transformer_options:
                 transformer_options["transformer_index"] += 1
-        elif isinstance(layer, Upsample):
-            x = layer(x, output_shape=output_shape)
+            x = rearrange(x, 'b c f h w -> (b f) c h w')
         else:
             x = layer(x)
     return x
@@ -52,19 +54,8 @@ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
     support it as an extra input.
     """
 
-    def forward(self, x, emb, context=None, batch_size=None):
-        for layer in self:
-            if isinstance(layer, TimestepBlock):
-                x = layer(x, emb, batch_size=batch_size)
-            elif isinstance(layer, SpatialTransformer):
-                x = layer(x, context)
-            elif isinstance(layer, TemporalTransformer):
-                x = rearrange(x, '(b f) c h w -> b c f h w', b=batch_size)
-                x = layer(x, context)
-                x = rearrange(x, 'b c f h w -> (b f) c h w')
-            else:
-                x = layer(x)
-        return x
+    def forward(self, *args, **kwargs):
+        return forward_timestep_embed(self, *args, **kwargs)
 
 class Downsample(nn.Module):
     """
@@ -340,6 +331,32 @@ def context_processor(context, t, img_emb=None, temporal_size=16, concat_only=Fa
 
     return context
 
+def apply_control(h, control, name, cond_idx=None):
+    if control is not None and name in control and len(control[name]) > 0:
+        frames = h.shape[0]
+        ctrl = control[name].pop()
+        if ctrl is not None:    
+            try:
+                if cond_idx is not None and ctrl.shape[0] > frames:
+                    ctrl_frames_list = list(range(ctrl.shape[0]))
+                    ctrl_frames = len(ctrl_frames_list)
+
+                    idxs = (
+                        ctrl_frames_list[ctrl_frames // 2:] if cond_idx == 0 else \
+                            ctrl_frames_list[:ctrl_frames // 2]
+                    )
+                    
+                    ctrl = ctrl[idxs]
+
+                h += ctrl
+            except Exception as e:
+                if h.shape != ctrl.shape:
+                    logging.warning(
+                        "warning control could not be applied {} {}".format(h.shape, ctrl.shape)
+                    )
+                logging.warning(e)
+    return h
+
 class UNetModel(nn.Module):
     """
     The full UNet model with attention and timestep embedding.
@@ -677,13 +694,16 @@ class UNetModel(nn.Module):
             features_adapter=None, 
             fs=None, 
             img_emb=None,
+            control=None,
             transformer_options={}, 
+            cond_idx=None,
             **kwargs
         ):
         
         if any([fs is None, img_emb is None, cc_concat is None]):
             raise ValueError("One or more of the required inputs for UNet Forward is None.")
         
+        cond_idx = transformer_options.get("cond_idx", None)
         transformer_options['original_shape'] = list(x.shape)
         transformer_options['transformer_index'] = 0
         transformer_patches = transformer_options.get("patches", {})
@@ -726,21 +746,73 @@ class UNetModel(nn.Module):
         hs = []
 
         for id, module in enumerate(self.input_blocks):
-            h = module(h, emb, context=context, batch_size=b)
+            transformer_options["block"] = ("input", id)
+            #h = module(h, emb, context=context, batch_size=b)
+            h = forward_timestep_embed(
+                module, 
+                h, 
+                emb, 
+                context=context, 
+                batch_size=b, 
+                transformer_options=transformer_options
+            )
+            h = apply_control(h, control, 'input', cond_idx)
+
+            if "input_block_patch" in transformer_patches:
+                patch = transformer_patches["input_block_patch"]
+                for p in patch:
+                    h = p(h, transformer_options)
+
             if id ==0 and self.addition_attention:
-                h = self.init_attn(h, emb, context=context, batch_size=b)
+                h = forward_timestep_embed(
+                self.init_attn, 
+                h, 
+                emb, 
+                context=context, 
+                batch_size=b, 
+                transformer_options=transformer_options
+            )
             ## plug-in adapter features
             if ((id+1)%3 == 0) and features_adapter is not None:
                 h = h + features_adapter[adapter_idx]
                 adapter_idx += 1
             hs.append(h)
+            if "input_block_patch_after_skip" in transformer_patches:
+                patch = transformer_patches["input_block_patch_after_skip"]
+                for p in patch:
+                    h = p(h, transformer_options)
         if features_adapter is not None:
             assert len(features_adapter)==adapter_idx, 'Wrong features_adapter'
+        transformer_options["block"] = ("middle", 0)
+        h = forward_timestep_embed(
+                self.middle_block, 
+                h, 
+                emb, 
+                context=context, 
+                batch_size=b, 
+                transformer_options=transformer_options
+            )
+        h = apply_control(h, control, 'middle', cond_idx)
+        for id, module in enumerate(self.output_blocks):
+            transformer_options["block"] = ("output", id)
+            hsp = hs.pop()
+            hsp = apply_control(hsp, control, 'output', cond_idx)
 
-        h = self.middle_block(h, emb, context=context, batch_size=b)
-        for module in self.output_blocks:
-            h = torch.cat([h, hs.pop()], dim=1)
-            h = module(h, emb, context=context, batch_size=b)
+            if "output_block_patch" in transformer_patches:
+                patch = transformer_patches["output_block_patch"]
+                for p in patch:
+                    h, hsp = p(h, hsp, transformer_options)
+
+            h = torch.cat([h, hsp], dim=1)
+            del hsp
+            h = forward_timestep_embed(
+                module, 
+                h, 
+                emb, 
+                context=context, 
+                batch_size=b, 
+                transformer_options=transformer_options
+            )
         h = h.type(x.dtype)
         h = self.out(h)
         

nodes.py

@@ -124,6 +124,7 @@ class DynamiCrafterProcessor:
         
         outs = []
         for i in range(batch_size):
+            model_in_kwargs['transformer_options']['cond_idx'] = i
             x_out = apply_model(
                     x_in[[i]], 
                     t=torch.cat([t[:1]]),