mimic-video 0.0.1__py3-none-any.whl

mimic_video/__init__.py

@@ -0,0 +1,2 @@
+
+from mimic_video.mimic_video import MimicVideo

mimic_video/mimic_video.py

@@ -0,0 +1,343 @@
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList, Linear
+
+import torch.nn.functional as F
+
+import einx
+from einops import einsum, rearrange
+from einops.layers.torch import Rearrange
+
+from x_mlps_pytorch import create_mlp
+
+from torch_einops_utils import (
+    pad_left_ndim,
+    align_dims_left
+)
+
+# ein notation
+
+# b - batch
+# h - heads
+# g - groups
+# n - sequence
+# i, j - sequence (source, target)
+# d - feature dimension
+
+# functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+# tensor function
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+def l2norm(t, eps = 1e-10):
+    return F.normalize(t, dim = -1, eps = eps)
+
+# time
+
+# they follow p0's research finding with the beta distribution
+# lets stick with 0 noise to 1 data instead of the reverse
+
+def default_sample_time_fn(time, s = 0.999):
+    return torch.sqrt(s - time)
+
+class RandomFourierEmbed(Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.proj = nn.Sequential(
+            Rearrange('... -> ... 1'),
+            nn.Linear(1, dim)
+        )
+
+        self.proj.requires_grad_(False)
+
+    def forward(self, times):
+        rand_proj = self.proj(times)
+        return torch.cos(2 * torch.pi * rand_proj)
+
+# adaptive rmsnorm
+
+class AdaptiveRMSNorm(Module):
+    def __init__(
+        self,
+        dim,
+        dim_time_cond,
+        eps = 1e-6
+    ):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.eps = eps
+
+        self.to_modulation = Linear(dim_time_cond, dim * 3, bias = False)
+        self.split_modulation = Rearrange('b (three d) -> three b 1 d', three = 3)
+
+        nn.init.zeros_(self.to_modulation.weight)
+
+    def forward(
+        self,
+        tokens,
+        time_cond
+    ):
+
+        if time_cond.ndim == 1:
+            time_cond = pad_left_ndim(time_cond, 1)
+
+        modulations = self.to_modulation(time_cond)
+
+        scale, shift, gate = self.split_modulation(modulations)
+
+        normed = l2norm(tokens, self.eps) * self.scale
+
+        adaptive_normed = normed * (scale + 1.) + shift
+
+        return adaptive_normed, gate
+
+# attention
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_context = None,
+        dim_head = 64,
+        heads = 8,
+        kv_heads = 2
+    ):
+        super().__init__()
+        dim_q_inner = dim_head * heads
+        dim_kv_inner = dim_head * kv_heads
+        dim_context = default(dim_context, dim)
+
+        self.scale = dim_head ** -0.5
+
+        self.to_queries = Linear(dim, dim_q_inner, bias = False)
+        self.to_keys_values = Linear(dim_context, dim_kv_inner * 2, bias = False)
+        self.to_out = Linear(dim_q_inner, dim, bias = False)
+
+        assert divisible_by(heads, kv_heads)
+        groups = heads // kv_heads
+
+        self.split_q_heads = Rearrange('b n (g h d) -> b g h n d', g = groups, d = dim_head)
+        self.split_kv_heads = Rearrange('b n (h d) -> b h n d', d = dim_head)
+        self.merge_heads = Rearrange('b g h n d -> b n (g h d)')
+
+    def forward(
+        self,
+        tokens,
+        context = None,
+        context_mask = None
+    ):
+        context = default(context, tokens)
+
+        queries = self.to_queries(tokens)
+        keys, values = self.to_keys_values(context).chunk(2, dim = -1)
+
+        queries = self.split_q_heads(queries)
+        keys, values = tuple(self.split_kv_heads(t) for t in (keys, values))
+
+        queries = queries * self.scale
+
+        sim = einsum(queries, keys, 'b g h i d, b h j d -> b g h i j')
+
+        if exists(context_mask):
+            mask_value = max_neg_value(sim)
+            sim = einx.where('b j, b g h i j,', context_mask, sim, mask_value)
+
+        attn = sim.softmax(dim = -1)
+
+        out = einsum(attn, values, 'b g h i j, b h j d -> b g h i d')
+
+        out = self.merge_heads(out)
+
+        return self.to_out(out)
+
+# feedforward
+
+class SwiGLUFeedForward(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        expansion_factor = 4.,
+    ):
+        super().__init__()
+        dim_inner = int(dim * expansion_factor * 2 / 3)
+
+        self.proj_in = nn.Linear(dim, dim_inner * 2)
+        self.proj_out = nn.Linear(dim_inner, dim)
+
+    def forward(
+        self,
+        tokens
+    ):
+        hidden, gates = self.proj_in(tokens).chunk(2, dim = -1)
+
+        out = hidden * F.gelu(gates)
+
+        return self.proj_out(out)
+
+# classes
+
+class MimicVideo(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_video_hidden,
+        dim_action = 20,
+        depth = 8,
+        dim_head = 64,
+        heads = 8,
+        expansion_factor = 4.,
+        dim_time_cond = None,
+        sample_time_fn = None
+    ):
+        super().__init__()
+
+        # flow related
+
+        self.sample_time_fn = default(sample_time_fn, default_sample_time_fn)
+
+        # embed
+
+        self.to_action_tokens = Linear(dim_action, dim)
+
+        dim_time_cond = default(dim_time_cond, dim * 2)
+
+        self.to_time_cond = nn.Sequential(
+            RandomFourierEmbed(dim),
+            create_mlp(dim_in = dim, dim = dim_time_cond, depth = 2, activation = nn.SiLU())
+        )
+
+        self.video_hidden_norm = nn.RMSNorm(dim_video_hidden)
+
+        # transformer
+
+        layers = []
+
+        for _ in range(depth):
+            attn_adanorm = AdaptiveRMSNorm(dim = dim, dim_time_cond = dim_time_cond)
+
+            attn = Attention(dim = dim, dim_head = dim_head, heads = heads)
+
+            cross_attn_adanorm = AdaptiveRMSNorm(dim = dim, dim_time_cond = dim_time_cond)
+
+            cross_attn = Attention(dim = dim, dim_head = dim_head, dim_context = dim_video_hidden, heads = heads)
+
+            ff_adanorm = AdaptiveRMSNorm(dim = dim, dim_time_cond = dim_time_cond)
+
+            ff = SwiGLUFeedForward(dim = dim, expansion_factor = expansion_factor)
+
+            layers.append(ModuleList([
+                attn_adanorm,
+                attn,
+                cross_attn_adanorm,
+                cross_attn,
+                ff_adanorm,
+                ff
+            ]))
+
+        self.layers = ModuleList(layers)
+
+        # predictions
+
+        self.to_pred_action_flow = nn.Sequential(
+            nn.RMSNorm(dim),
+            Linear(dim, dim_action)
+        )
+
+    def forward(
+        self,
+        actions,
+        video_hiddens, # they use layer 19 of cosmos predict, at first denoising step. that's all
+        *,
+        time = None,
+        context_mask = None,
+    ):
+
+        is_training = not exists(time)
+
+        # handle flow time conditioning
+
+        if is_training:
+            batch, device = actions.shape[0], actions.device
+
+            time = torch.rand((batch,), device = device)
+            time = self.sample_time_fn(time)
+
+            noise = torch.randn_like(actions)
+            flow = actions - noise
+
+            actions, left_aligned_time = align_dims_left((actions, time))
+
+            noised = noise.lerp(actions, left_aligned_time)
+        else:
+            noised = actions
+
+        time_cond = self.to_time_cond(time)
+
+        # handle video hiddens
+
+        video_hiddens = self.video_hidden_norm(video_hiddens)
+
+        # embed
+
+        tokens = self.to_action_tokens(noised)
+
+        # transformer layers
+
+        for (
+            attn_norm,
+            attn,
+            cross_attn_norm,
+            cross_attn,
+            ff_norm,
+            ff
+        ) in self.layers:
+
+            # cross attention
+
+            residual = tokens
+
+            tokens, gate = cross_attn_norm(tokens, time_cond)
+
+            tokens = residual + cross_attn(tokens, context = video_hiddens, context_mask = context_mask) * gate
+
+            # self attention
+
+            residual = tokens
+
+            tokens, gate = attn_norm(tokens, time_cond)
+
+            tokens = residual + attn(tokens) * gate
+
+            # feedforward
+
+            residual = tokens
+
+            tokens, gate = ff_norm(tokens, time_cond)
+
+            tokens = residual + ff(tokens) * gate
+
+        # prediction
+
+        pred_flow = self.to_pred_action_flow(tokens)
+
+        if not is_training:
+            return pred_flow
+
+        # mse flow loss
+
+        flow_loss = F.mse_loss(pred_flow, flow)
+        return flow_loss

mimic_video-0.0.1.dist-info/METADATA

@@ -0,0 +1,78 @@
+Metadata-Version: 2.4
+Name: mimic-video
+Version: 0.0.1
+Summary: Mimic Video
+Project-URL: Homepage, https://pypi.org/project/mimic-video/
+Project-URL: Repository, https://github.com/lucidrains/mimic-video
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2025 Phil Wang
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: artificial intelligence,attention mechanism,deep learning,video language action model
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: einops>=0.8.1
+Requires-Dist: einx>=0.3.0
+Requires-Dist: torch-einops-utils>=0.0.8
+Requires-Dist: torch>=2.5
+Requires-Dist: x-mlps-pytorch
+Provides-Extra: examples
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Description-Content-Type: text/markdown
+
+<img src="./mimic-video.png" width="450px"></img>
+
+## Mimic Video (wip)
+
+Implementation of [Mimic-Video](https://mimic-video.github.io/), Video-Action Models for Generalizable Robot Control Beyond VLAs
+
+## Contributing
+
+First make sure `pytest` and test dependencies are installed with
+
+```shell
+$ pip install '.[test]'
+```
+
+Then add your test to `tests/test_mimic_video.py` and run
+
+```shell
+$ pytest tests
+```
+
+That's it
+
+## Citations
+
+```bibtex
+@inproceedings{Pai2025mimicvideoVM,
+    title   = {mimic-video: Video-Action Models for Generalizable Robot Control Beyond VLAs},
+    author  = {Jonas Pai and Liam Achenbach and Victoriano Montesinos and Benedek Forrai and Oier Mees and Elvis Nava},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:283920528}
+}
+```

mimic_video-0.0.1.dist-info/RECORD

@@ -0,0 +1,6 @@
+mimic_video/__init__.py,sha256=-4HP_pbT4YLhRUwNwuL4qyLHbgDyQ099nHL7eVi0_Ag,48
+mimic_video/mimic_video.py,sha256=aejvjr1F3A7pZFikf-kEgeOpi1_53xVddBMpDPoxA90,8272
+mimic_video-0.0.1.dist-info/METADATA,sha256=414y344JcuIKQJss7d9riTrHszIwthHW8DDSSuRntdo,2960
+mimic_video-0.0.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mimic_video-0.0.1.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+mimic_video-0.0.1.dist-info/RECORD,,

mimic_video-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

mimic_video-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Phil Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.