Merge pull request #6 from AlexKoff88/ak/fixes

Fixes

tomeov/import_utils.py

@@ -1,21 +1,23 @@
 def is_diffusers_available():
     try:
-        import stable_diffusion
-        import utils
+        import diffusers
         return True
     except ImportError:
+        print("diffusers library is not available. Please install it to use Token Merging.")
         return False
 
 def is_openclip_available():
     try:
-        import open_clip_torch
+        import open_clip
         return True
     except ImportError:
+        print("OpenCLIP library is not available. Please install it to use Token Merging.")
         return False
 
 def is_timm_available():
     try:
         import timm
         return True
     except ImportError:
+        print("Timm library is not available. Please install it to use Token Merging.")
         return False

tomeov/stable_diffusion.py

@@ -4,7 +4,7 @@
 
 from . import merge
 from .utils import isinstance_str
-from .merge import merge_wavg
+
 
 
 def compute_merge(x: torch.Tensor, tome_info: Dict[str, Any]) -> Tuple[Callable, ...]:
@@ -54,6 +54,22 @@ def _forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tenso
     return ToMeBlock
 
 
+def merge_wavg(
+    merge: Callable, x: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Applies the merge function by taking a weighted average based on token size.
+    Returns the merged tensor and the new token sizes.
+    """
+    size = torch.ones_like(x[..., 0, None])
+
+    x = merge(x * size, mode="sum")
+    size = merge(size, mode="sum")
+
+    x = x / size
+    return x
+
+
 def make_diffusers_tome_block(block_class: Type[torch.nn.Module]) -> Type[torch.nn.Module]:
     """
     Make a patched class for a diffusers model.