reference unet

utils/module_utils.py

@@ -1,6 +1,7 @@
 import torch
 
 from ..modules.attn_wrapper import get_attn_wrapper
+from ..modules.unet_wrapper import get_unet_wrapper
 from ..modules.ref_controller import RefController
 
 
@@ -70,4 +71,13 @@ def setup_ref_attn(model, attn_type):
     for _, transformer_module in temp_transformers:
         ref_controller.add_temporal_transformer(transformer_module)
 
-    return ref_controller
+    return ref_controller
+
+
+def setup_ref_unet(model, attn_type):
+    if not hasattr(model.model.diffusion_model, 'is_ref'):
+        ref_controller = setup_ref_attn(model, attn_type)
+        model.model.diffusion_model.__class__ = get_unet_wrapper(model.model.diffusion_model.__class__, ref_controller)
+        return ref_controller
+    return None
+