mimic-video 0.0.31__py3-none-any.whl

mimic_video/__init__.py

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

mimic_video/cosmos_predict.py

@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import torch
+from torch.nn import Module
+from torch import nn, Tensor
+from einops import rearrange
+
+from diffusers.models.transformers.transformer_cosmos import CosmosTransformer3DModel
+from diffusers.models.autoencoders.autoencoder_kl_cosmos import AutoencoderKLCosmos
+from diffusers.schedulers.scheduling_edm_euler import EDMEulerScheduler
+from transformers import T5EncoderModel, T5TokenizerFast, T5Config
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def identity(t):
+    return t
+
+def default(v, d):
+    return v if exists(v) else d
+
+# constants
+
+TINY_TRANSFORMER_CONFIG = dict(
+    in_channels = 16,
+    out_channels = 16,
+    num_attention_heads = 1,
+    attention_head_dim = 16,
+    mlp_ratio = 1.0,
+    text_embed_dim = 32,
+    adaln_lora_dim = 32,
+    patch_size = (1, 2, 2),
+    max_size = (4, 16, 16),
+    extra_pos_embed_type = None,
+    concat_padding_mask = False,
+)
+
+TINY_VAE_CONFIG = dict(
+    in_channels = 3,
+    out_channels = 3,
+    latent_channels = 16,
+    encoder_block_out_channels = (8, 16),
+    decode_block_out_channels = (8, 16),
+    temporal_compression_ratio = 4,
+    spatial_compression_ratio = 4,
+    num_layers = 1,
+    attention_resolutions = (),
+    resolution = 64,
+)
+
+TINY_T5_CONFIG = dict(
+    vocab_size = 32128,
+    d_model = 32,
+    d_kv = 8,
+    d_ff = 64,
+    num_layers = 1,
+    num_heads = 1,
+)
+
+REAL_TRANSFORMER_CONFIG = dict(
+    in_channels = 16,
+    out_channels = 16,
+    num_attention_heads = 32,
+    attention_head_dim = 128,
+    mlp_ratio = 4.0,
+    text_embed_dim = 1024,
+    patch_size = (1, 2, 2),
+    max_size = (128, 240, 240),
+    extra_pos_embed_type = "learnable",
+    concat_padding_mask = True,
+)
+
+REAL_VAE_CONFIG = dict(
+    in_channels = 3,
+    out_channels = 3,
+    latent_channels = 16,
+    encoder_block_out_channels = (128, 256, 512, 512),
+    decode_block_out_channels = (256, 512, 512, 512),
+    temporal_compression_ratio = 8,
+    spatial_compression_ratio = 8,
+)
+
+REAL_T5_CONFIG = dict(
+    vocab_size = 32128,
+    d_model = 1024,
+    d_kv = 64,
+    d_ff = 2048,
+    num_layers = 12,
+    num_heads = 16,
+)
+
+# main class
+
+class CosmosPredictWrapper(Module):
+    """
+    Wraps Cosmos VAE + DiT for extracting hidden states from a video.
+    Supports proper EDM Euler denoising steps.
+    """
+    
+    def __init__(
+        self,
+        model_name: str = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
+        extract_layers: int | list[int] | None = None,
+        random_weights: bool = False,
+        tiny: bool = False,
+        normalize = lambda t: (t - 0.5) * 2.0,
+        extract_layer: int | None = None
+    ):
+        super().__init__()
+        extract_layers = default(extract_layers, extract_layer)
+        extract_layers = default(extract_layers, 19)
+
+        self.extract_layers = [extract_layers] if isinstance(extract_layers, int) else extract_layers
+        self.return_list = isinstance(extract_layers, list)
+
+        self.hook_handles: list = []
+        self.cached_hidden_states: list[Tensor] = []
+        
+        if random_weights:
+            self._init_random_weights(tiny = tiny)
+        else:
+            self._init_pretrained(model_name)
+        
+        # Initialize scheduler
+        self.scheduler = EDMEulerScheduler()
+        
+        # store hidden dim for consumers
+        self.dim_latent = self.transformer.config.num_attention_heads * self.transformer.config.attention_head_dim
+
+        # maybe normalize
+        self.normalize = normalize
+
+        self._register_hook()
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def _init_pretrained(self, model_name: str):
+        """Load pretrained weights from HuggingFace"""
+        from diffusers import CosmosVideoToWorldPipeline
+        
+        pipeline = CosmosVideoToWorldPipeline.from_pretrained(model_name)
+        
+        # Extract components we need
+        self.vae = pipeline.vae
+        self.transformer = pipeline.transformer
+        self.text_encoder = pipeline.text_encoder
+        self.tokenizer = pipeline.tokenizer
+
+        # Clean up pipeline
+        del pipeline
+
+    def _init_random_weights(self, tiny: bool = False):
+        """Initialize with random weights for testing"""
+        
+        transformer_config = TINY_TRANSFORMER_CONFIG if tiny else REAL_TRANSFORMER_CONFIG
+        vae_config = TINY_VAE_CONFIG if tiny else REAL_VAE_CONFIG
+        t5_config_dict = TINY_T5_CONFIG if tiny else REAL_T5_CONFIG
+
+        num_layers = max(2, *[layer + 1 for layer in self.extract_layers])
+        if not tiny:
+            num_layers = max(28, num_layers)
+        
+        self.transformer = CosmosTransformer3DModel(
+            num_layers = num_layers,
+            **transformer_config
+        )
+        
+        self.vae = AutoencoderKLCosmos(**vae_config)
+        
+        t5_config = T5Config(**t5_config_dict)
+        self.text_encoder = T5EncoderModel(t5_config)
+        self.tokenizer = T5TokenizerFast.from_pretrained("google-t5/t5-small")
+
+    def __del__(self):
+        if not hasattr(self, 'hook_handles'):
+            return
+
+        for handle in self.hook_handles:
+            handle.remove()
+
+    def _register_hook(self):
+        assert hasattr(self.transformer, 'transformer_blocks'), 'transformer must have transformer_blocks'
+        
+        for layer_index in self.extract_layers:
+            assert len(self.transformer.transformer_blocks) > layer_index, f'layer {layer_index} out of bounds'
+
+            target_layer = self.transformer.transformer_blocks[layer_index]
+
+            def hook_fn(module, inp, out):
+                self.cached_hidden_states.append(out.detach().cpu())
+
+            handle = target_layer.register_forward_hook(hook_fn)
+            self.hook_handles.append(handle)
+
+    def forward(
+        self,
+        videos: Tensor,
+        prompts: str | list[str] | None = None,
+        prompt_token_ids: Tensor | None = None,
+        num_inference_steps: int = 1,
+    ) -> Tensor:
+        """
+        videos: (batch, frames, channels, height, width) in [0, 1]
+        num_inference_steps: number of denoising steps to run
+        returns: hidden states tensor from the specified transformer layer (from first step)
+        """
+        batch, t, c, h, w = videos.shape
+
+        assert exists(prompts) ^ exists(prompt_token_ids)
+
+        # Scale videos from [0, 1] to [-1, 1] for Cosmos VAE
+
+        videos = self.normalize(videos)
+
+        if isinstance(prompts, str):
+            prompts = [prompts] * batch
+
+        self.cached_hidden_states.clear()
+        
+        # Move video to device and rearrange for VAE: (B, T, C, H, W) -> (B, C, T, H, W)
+        videos = rearrange(videos, 'b t c h w -> b c t h w')
+        
+        with torch.inference_mode():
+            # 1. encode video to latents via VAE
+
+            latents = self.vae.encode(videos).latent_dist.sample()
+            
+            # 2. maybe encode text prompts
+
+            if exists(prompt_token_ids):
+                text_inputs = dict(input_ids = prompt_token_ids)
+            else:
+                text_inputs = self.tokenizer(
+                    prompts, 
+                    return_tensors = "pt",
+                    padding = True,
+                    truncation = True,
+                    max_length = 512
+                )
+
+            encoder_hidden_states = self.text_encoder(**text_inputs).last_hidden_state
+            
+            # 3. Setup scheduler timesteps
+            self.scheduler.set_timesteps(num_inference_steps, device = self.device)
+            timesteps = self.scheduler.timesteps
+            
+            # 4. Add noise to latents (start from pure noise scaled by initial sigma)
+            noise = torch.randn_like(latents)
+            latents = latents + noise * self.scheduler.init_noise_sigma
+            
+            # 5. Denoising loop
+            for i, timestep in enumerate(timesteps):
+                # Scale model input
+                latent_model_input = self.scheduler.scale_model_input(latents, timestep)
+                
+                # Predict noise residual
+                noise_pred = self.transformer(
+                    hidden_states = latent_model_input,
+                    encoder_hidden_states = encoder_hidden_states,
+                    timestep = timestep.expand(batch),
+                    return_dict = False
+                )[0]
+                
+                # Compute previous noisy sample
+                latents = self.scheduler.step(noise_pred, timestep, latents, return_dict = False)[0]
+
+        assert len(self.cached_hidden_states) >= len(self.extract_layers), 'hidden states not captured'
+        
+        # Return hidden states from the first denoising step
+        hiddens = self.cached_hidden_states[:len(self.extract_layers)]
+
+        for hidden in hiddens:
+            assert hidden.shape[-1] == self.dim_latent, f'hidden dim mismatch: expected {self.dim_latent}, got {hidden.shape[-1]}'
+
+        if not self.return_list:
+            return hiddens[0]
+
+        return hiddens

mimic_video/mimic_video.py

@@ -0,0 +1,797 @@
+from __future__ import annotations
+from functools import partial
+
+import torch
+from torch import nn, cat, stack, is_tensor, tensor
+from torch.nn import Module, ModuleList, Linear, GRU
+
+import torch.nn.functional as F
+
+import einx
+from einops import einsum, rearrange, repeat, reduce
+from einops.layers.torch import Rearrange
+
+from x_mlps_pytorch import create_mlp
+
+from tqdm import tqdm
+
+from torch_einops_utils import (
+    lens_to_mask,
+    pad_left_ndim,
+    align_dims_left,
+    pad_at_dim,
+    pack_with_inverse,
+    masked_mean,
+    tree_map_tensor
+)
+
+from hyper_connections.mHCv2 import get_init_and_expand_reduce_stream_functions
+
+# ein notation
+
+# b - batch
+# h - heads
+# g - groups
+# n - sequence
+# i, j - sequence (source, target)
+# d - feature dimension
+
+# constants
+
+LinearNoBias = partial(Linear, bias = False)
+
+# functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def identity(t):
+    return t
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+# wrappers
+
+def eval_no_grad(fn):
+    def inner(*args, **kwargs):
+        with torch.no_grad():
+            fn.eval()
+            return fn(*args, **kwargs)
+
+    return inner
+
+# tensor function
+
+def cast_tensor(val, device = None):
+    return tensor(val, device = device) if not is_tensor(val) else val
+
+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)
+
+# token shift from Peng et al. of RWKV
+# cheap way to generate relative positions
+
+def shift_feature_dim(t):
+    x, x_shift = t.chunk(2, dim = -1)
+    x_shift = pad_at_dim(x_shift, (1, -1), dim = 1)
+    return cat((x, x_shift), dim = -1)
+
+# action normalization
+
+class Normalizer(Module):
+    def __init__(
+        self,
+        mean,
+        std,
+        eps = 1e-6
+    ):
+        super().__init__()
+        assert (std > 0.).all(), 'std must be positive'
+        self.eps = eps
+
+        self.register_buffer('mean', mean)
+        self.register_buffer('std', std)
+
+    def normalize(self, t):
+        mean, std = self.mean, self.std
+        return (t - mean) / std.clamp_min(self.eps)
+
+    def inverse_normalize(self, t):
+        mean, std = self.mean, self.std
+        return (t * std) + mean
+
+# 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,
+        ada_ln_zero_bias = -5.
+    ):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.eps = eps
+
+        self.to_modulation = LinearNoBias(dim_time_cond, dim * 3)
+        self.split_modulation = Rearrange('... (three d) -> three ... d', three = 3)
+
+        nn.init.zeros_(self.to_modulation.weight)
+
+        self.ada_ln_zero_bias = ada_ln_zero_bias
+
+    def forward(
+        self,
+        tokens,
+        time_cond
+    ):
+        if time_cond.ndim == 2:
+            time_cond = rearrange(time_cond, 'b d -> b 1 d')
+
+        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
+
+        gate_with_bias = (gate + self.ada_ln_zero_bias).sigmoid()
+
+        return adaptive_normed, gate_with_bias
+
+# attention
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_context = None,
+        dim_head = 64,
+        heads = 8,
+        kv_heads = 2,
+        attn_gate_value = True,
+        norm_context = False
+    ):
+        super().__init__()
+        dim_q_inner = dim_head * heads
+        dim_kv_inner = dim_head * kv_heads
+
+        dim_context = default(dim_context, dim)
+        self.context_norm = nn.RMSNorm(dim_context) if norm_context else nn.Identity()
+
+        self.scale = dim_head ** -0.5
+
+        self.to_queries = LinearNoBias(dim, dim_q_inner)
+        self.to_keys_values = LinearNoBias(dim_context, dim_kv_inner * 2)
+
+        self.attn_gate_value = nn.Sequential(LinearNoBias(dim, heads), Rearrange('b n (g h) -> b g h n 1', h = kv_heads))
+
+        self.to_out = LinearNoBias(dim_q_inner, dim)
+
+        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,
+        kv = None,
+        return_kv = False
+    ):
+        context = default(context, tokens)
+
+        queries = self.to_queries(tokens)
+        queries = self.split_q_heads(queries)
+
+        if not exists(kv):
+            context = self.context_norm(context)
+
+            keys, values = self.to_keys_values(context).chunk(2, dim = -1)
+            keys, values = tuple(self.split_kv_heads(t) for t in (keys, values))
+        else:
+            keys, values = kv
+
+        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')
+
+        # https://openreview.net/forum?id=1b7whO4SfY - should become standard practice
+
+        out = out * self.attn_gate_value(tokens).sigmoid()
+
+        out = self.merge_heads(out)
+
+        out = self.to_out(out)
+
+        if not return_kv:
+            return out
+
+        return out, stack((keys, values))
+
+# 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,
+        video_predict_wrapper: Module | None = None,
+        *,
+        dim_video_hidden = None,
+        action_chunk_len = 32,
+        dim_action = 20,
+        dim_joint_state = 32,
+        proprio_mask_prob = 0.1,
+        depth = 8,
+        dim_head = 64,
+        heads = 8,
+        expansion_factor = 4.,
+        ada_ln_zero_bias = -5.,
+        dim_time_cond = None,
+        sample_time_fn = None,
+        train_time_rtc = False,
+        train_time_rtc_max_delay = None,
+        num_residual_streams = 1,
+        mhc_kwargs: dict = dict(),
+        action_mean_std: Tensor | None = None,
+        joint_mean_std: Tensor | None = None,
+        num_task_ids = 0,
+        num_advantage_ids = 0,
+        advantage_cfg_dropout = 0.25,
+        extracted_video_layer_indices: list[int] | None = None
+    ):
+        super().__init__()
+
+        self.depth = depth
+
+        # maybe video predict
+
+        self.video_predict_wrapper = video_predict_wrapper
+
+        # action related
+
+        self.action_chunk_len = action_chunk_len
+        self.dim_action = dim_action
+
+        self.action_shape = (action_chunk_len, dim_action)
+
+        self.action_normalizer = None
+
+        if exists(action_mean_std):
+            assert action_mean_std.shape == (2, dim_action), f'must be in shape of (2 action_dim)'
+            self.action_normalizer = Normalizer(*action_mean_std)
+
+        # joint dim
+
+        self.dim_joint_state = dim_joint_state
+
+        dim_video_hidden = default(dim_video_hidden, video_predict_wrapper.dim_latent if exists(video_predict_wrapper) else None)
+
+        assert exists(dim_video_hidden), f'`dim_video_hidden` must be set or `video_predict_wrapper` passed in with `dim_latent`'
+
+        self.dim_video_hidden = dim_video_hidden
+
+        self.joint_normalizer = None
+
+        if exists(joint_mean_std):
+            assert joint_mean_std == (2, dim_joint_state)
+            self.joint_normalizer = Normalizer(*joint_mean_std)
+
+        # 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_fourier_embed = RandomFourierEmbed(dim) # used by deepmind, its fine
+        self.to_time_cond = create_mlp(dim_in = dim * 2, dim = dim_time_cond, depth = 2, activation = nn.SiLU())
+
+        # joint token related
+
+        self.to_joint_state_token = Linear(dim_joint_state, dim)
+
+        self.proprio_mask_prob = proprio_mask_prob
+        self.has_proprio_masking = proprio_mask_prob > 0.
+
+        self.proprio_mask_token = nn.Parameter(torch.randn(dim))
+
+        # video norm
+
+        self.video_hidden_norm = nn.RMSNorm(dim_video_hidden)
+
+        # manifold constrained hyper connections (mHC) from bytedance + deepseek
+
+        init_hyper_conn, self.expand_stream, self.reduce_stream = get_init_and_expand_reduce_stream_functions(num_residual_streams, dim = dim, add_stream_embed = True, **mhc_kwargs)
+
+        # rnn
+
+        self.rnn = GRU(dim, dim)
+
+        # 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, norm_context = True)
+
+            ff_adanorm = AdaptiveRMSNorm(dim = dim, dim_time_cond = dim_time_cond, ada_ln_zero_bias = ada_ln_zero_bias)
+
+            ff = SwiGLUFeedForward(dim = dim, expansion_factor = expansion_factor)
+
+            # maybe hyper connect
+
+            attn_residual = init_hyper_conn()
+            cross_attn_residual = init_hyper_conn()
+            ff_residual = init_hyper_conn()
+
+            layers.append(ModuleList([
+                cross_attn_residual,
+                cross_attn_adanorm,
+                cross_attn,
+                attn_residual,
+                attn_adanorm,
+                attn,
+                ff_residual,
+                ff_adanorm,
+                ff
+            ]))
+
+        self.layers = ModuleList(layers)
+
+        # predictions
+
+        self.to_pred_action_flow = nn.Sequential(
+            nn.RMSNorm(dim),
+            Linear(dim, dim_action, bias = False)
+        )
+
+        # inference related
+
+        # train time RTC related - https://arxiv.org/abs/2512.05964
+
+        self.train_time_rtc = train_time_rtc
+
+        assert not train_time_rtc or exists(train_time_rtc_max_delay)
+        self.train_time_rtc_max_delay = train_time_rtc_max_delay
+
+        # condition related
+
+        self.task_embed = nn.Embedding(num_task_ids, dim) if num_task_ids > 0 else None
+
+        self.advantage_embed = nn.Embedding(num_advantage_ids + 1, dim) if num_advantage_ids > 0 else None
+
+        assert advantage_cfg_dropout > 0.
+
+        self.advantage_cfg_dropout = advantage_cfg_dropout
+
+        # allow for researchers to explore beyond just one layer of pretrained
+        # we should also open up research into multiple pretrained models eventually
+
+        self.extracted_video_layer_indices = default(extracted_video_layer_indices, (0,) * depth)
+        assert len(self.extracted_video_layer_indices) == depth
+
+        # aux loss and device
+
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
+    # only action parameters
+
+    def action_parameters(self):
+        video_model_params = set(self.video_predict_wrapper.parameters()) if exists(self.video_predict_wrapper) else {}
+        return set(self.parameters()) - video_model_params
+
+    @property
+    def device(self):
+        return self.zero.device
+
+    @torch.no_grad()
+    def sample(
+        self,
+        steps = 16,
+        batch_size = 1,
+        prefix_action_chunk = None,
+        disable_progress_bar = False,
+        **kwargs
+    ):
+
+        self.eval()
+
+        inpainting = exists(prefix_action_chunk)
+
+        # times
+
+        times = torch.linspace(0., 1., steps + 1, device = self.device)[:-1]
+        delta = 1. / steps
+
+        # inpaint
+
+        if inpainting:
+            prefix_len = prefix_action_chunk.shape[1]
+            assert prefix_len < self.action_chunk_len
+
+            maybe_normed_prefix = prefix_action_chunk
+
+            if exists(self.action_normalizer):
+                maybe_normed_prefix = self.action_normalizer.normalize(prefix_action_chunk)
+
+            times = repeat(times, 'steps -> steps b n', b = batch_size, n = self.action_chunk_len).clone()
+            times[..., :prefix_len] = 1.
+
+        # noise
+
+        noise = torch.randn((batch_size, *self.action_shape), device = self.device)
+
+        # denoised action starts as noise
+
+        denoised = noise
+
+        cache = None
+
+        # denoise
+
+        for time in tqdm(times, disable = disable_progress_bar):
+
+            if inpainting:
+                denoised[:, :prefix_len] = maybe_normed_prefix
+
+            pred_flow, cache = self.forward(actions = denoised, time = time, cache = cache, return_cache = True, **kwargs)
+
+            denoised = denoised + delta * pred_flow
+
+        # handle action inverse norm
+
+        if exists(self.action_normalizer):
+            denoised = self.action_normalizer.inverse_normalize(denoised)
+
+        # final set, with unnormalized prefix, if inpainting
+
+        if inpainting:
+            denoised[:, :prefix_len] = prefix_action_chunk
+
+        return denoised
+
+    def forward(
+        self,
+        *,
+        actions,                    # (b na d)
+        joint_state,                # (b)
+        task_ids = None,            # (b)
+        advantage_ids = None,       # (b)
+        video = None,               # (b c t h w)
+        video_hiddens = None,       # (b nv dv) - they use layer 19 of cosmos predict, at first denoising step. that's all
+        context_mask = None,
+        time = None,                # () | (b) | (b n)
+        time_video_denoise = 0.,    # 0 is noise in the scheme i prefer - default to their optimal choice, but can be changed
+        prompts: list[str] | None = None,
+        prompt_token_ids = None,
+        detach_video_hiddens = False,
+        no_grad_video_model_forward = False,
+        cache = None,
+        return_cache = False,
+        return_flow = False
+    ):
+        assert not exists(self.video_predict_wrapper) or (exists(prompts) ^ exists(prompt_token_ids))
+        assert actions.shape[-2:] == self.action_shape
+
+        if exists(self.action_normalizer):
+            actions = self.action_normalizer.normalize(actions)
+
+        batch, device = actions.shape[0], actions.device
+        orig_actions = actions
+
+        is_training = not exists(time) and not return_flow
+
+        if not exists(cache):
+            # handle maybe extraction of video hiddens
+            # only if cache is not given
+
+            assert exists(video) ^ exists(video_hiddens)
+
+            if not exists(video_hiddens):
+                assert exists(self.video_predict_wrapper), f'`video_predict_wrapper` must be passed in if raw video is passed into MimicVideo'
+
+                video_forward_wrap = eval_no_grad if no_grad_video_model_forward else identity
+
+                video_hiddens = video_forward_wrap(self.video_predict_wrapper)(video, prompts = prompts, prompt_token_ids = prompt_token_ids)
+
+                video_hiddens = tree_map_tensor(lambda t: t.to(self.device).float(), video_hiddens) # maybe bfloat to float32
+
+                video_hiddens = tree_map_tensor(lambda t: pack_with_inverse(t, 'b * d')[0], video_hiddens)
+
+            # handle video hiddens
+
+            if detach_video_hiddens:
+                video_hiddens = video_hiddens.detach()
+
+            if not isinstance(video_hiddens, list):
+                video_hiddens = [video_hiddens]
+
+        # handle caching
+
+        prev_cached_video_hiddens_kv = cache if exists(cache) else ((None,) * self.depth)
+        next_cached_video_hiddens_kv = []
+
+        # handle flow time conditioning
+
+        if is_training:
+            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
+
+        # maybe train time rtc
+
+        action_loss_mask = None
+
+        if is_training and self.train_time_rtc:
+
+            rand_prefix_len = torch.randint(0, self.train_time_rtc_max_delay, (batch,), device = device)
+            action_prefix_mask = lens_to_mask(rand_prefix_len, self.action_chunk_len)
+
+            actions = einx.where('b na, b na d, b na d', action_prefix_mask, orig_actions, actions)
+            time = einx.where('b na, , b', action_prefix_mask, 1., time)
+
+            action_loss_mask = ~action_prefix_mask
+
+        if time.ndim == 0:
+            time = repeat(time, '-> b', b = batch)
+
+        # handle the video denoising times
+
+        time_video_denoise = cast_tensor(time_video_denoise)
+
+        if time_video_denoise.ndim == 0:
+            time_video_denoise = rearrange(time_video_denoise, '-> 1')
+
+        if time_video_denoise.shape[0] != batch:
+            time_video_denoise = repeat(time_video_denoise, '1 -> b', b = batch)
+
+        if time.ndim == 2:
+            time_video_denoise = repeat(time_video_denoise, 'b -> b n', n = time.shape[-1])
+
+        times = stack((time, time_video_denoise), dim = -1)
+
+        # embed
+
+        tokens = self.to_action_tokens(noised)
+
+        # setup empty tokens for various packed condition tokens
+
+        empty_token = tokens[:, 0:0]
+
+        # one layer of rnn for actions
+
+        rnn_out, _, = self.rnn(tokens)
+        tokens = rnn_out + tokens
+
+        #  mask joint state token for proprioception masking training
+
+        if exists(self.joint_normalizer):
+            joint_state = self.joint_normalizer.normalize(joint_state)
+
+        joint_state_token = self.to_joint_state_token(joint_state)
+
+        if self.training and self.has_proprio_masking:
+            mask = torch.rand((batch,), device = device) < self.proprio_mask_prob
+
+            joint_state_token = einx.where('b, d, b d', mask, self.proprio_mask_token, joint_state_token)
+
+        # setup task
+
+        task_embed = empty_token
+
+        if exists(task_ids):
+            assert exists(self.task_embed)
+            task_embed = self.task_embed(task_ids)
+
+        # setup maybe advantage
+
+        advantage_embed = empty_token
+
+        if exists(advantage_ids):
+            assert exists(self.advantage_embed)
+
+            advantage_ids = advantage_ids + 1 # 0 for dropout
+            cfg_dropout = torch.rand_like(advantage_ids.float()) < self.advantage_cfg_dropout
+
+            advantage_ids = einx.where('b, , b', cfg_dropout, 0, advantage_ids)
+
+            advantage_embed = self.advantage_embed(advantage_ids)
+
+        # determine time - need to handle the sequence dimension given train time RTC and various conditioning tokens
+
+        if times.ndim == 3:
+            joint_task_advantage_times = 1 + int(exists(advantage_ids)) + int(exists(task_ids))
+
+            times = pad_at_dim(times, (joint_task_advantage_times, 0), dim = 1, value = 1.) # handle joint state token on the action
+
+        # fourier embed and mlp to time condition
+
+        fourier_embed = self.to_fourier_embed(times)
+
+        fourier_embed = rearrange(fourier_embed, '... times d -> ... (times d)')
+
+        time_cond = self.to_time_cond(fourier_embed)
+
+        # pack with action tokens for attention tower
+
+        tokens, inverse_pack = pack_with_inverse((advantage_embed, task_embed, joint_state_token, tokens), 'b * d')
+
+        # maybe expand streams
+
+        tokens = self.expand_stream(tokens)
+
+        # transformer layers
+
+        for ((
+            maybe_cross_attn_mhc,
+            cross_attn_norm,
+            cross_attn,
+            maybe_attn_mhc,
+            attn_norm,
+            attn,
+            maybe_ff_mhc,
+            ff_norm,
+            ff
+        ), layer_video_hidden_index, cached_video_kv) in zip(self.layers, self.extracted_video_layer_indices, prev_cached_video_hiddens_kv):
+
+            # cross attention
+
+            tokens, add_residual = maybe_cross_attn_mhc(tokens)
+
+            tokens, gate = cross_attn_norm(tokens, time_cond)
+
+            layer_video_hidden = None
+
+            if exists(video_hiddens):
+                layer_video_hidden = video_hiddens[layer_video_hidden_index]
+
+            cross_attn_out, video_kv = cross_attn(tokens, context = layer_video_hidden, context_mask = context_mask, kv = cached_video_kv, return_kv = True)
+
+            tokens = add_residual(cross_attn_out * gate)
+
+            if return_cache:
+                next_cached_video_hiddens_kv.append(video_kv)
+
+            # self attention
+
+            tokens, add_residual = maybe_attn_mhc(tokens)
+
+            tokens, gate = attn_norm(tokens, time_cond)
+
+            tokens = add_residual(attn(tokens) * gate)
+
+            # prepare feedforward
+
+            tokens, add_residual = maybe_ff_mhc(tokens)
+
+            tokens, gate = ff_norm(tokens, time_cond)
+
+            # shift along time for action tokens for cheap relative positioning, which is better than messing with rope with such short action chunks
+
+            *non_action_tokens, tokens = inverse_pack(tokens)
+
+            tokens = shift_feature_dim(tokens)
+
+            tokens, _ = pack_with_inverse((*non_action_tokens, tokens), 'b * d')
+
+            # feedforward
+
+            tokens = add_residual(ff(tokens) * gate)
+
+        # maybe reduce streams
+
+        tokens = self.reduce_stream(tokens)
+
+        # remove joint token
+
+        *_, tokens = inverse_pack(tokens)
+
+        # prediction
+
+        pred_flow = self.to_pred_action_flow(tokens)
+
+        if not is_training:
+            # flow
+
+            out = pred_flow
+        else:
+            # mse flow loss
+
+            flow_loss = F.mse_loss(pred_flow, flow, reduction = 'none')
+
+            out = masked_mean(flow_loss, action_loss_mask)
+
+        if not return_cache:
+            return out
+
+        # handle returning of cache
+
+        return out, next_cached_video_hiddens_kv

mimic_video-0.0.31.dist-info/METADATA

@@ -0,0 +1,174 @@
+Metadata-Version: 2.4
+Name: mimic-video
+Version: 0.0.31
+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: hyper-connections>=0.4.3
+Requires-Dist: torch-einops-utils>=0.0.14
+Requires-Dist: torch>=2.5
+Requires-Dist: tqdm
+Requires-Dist: x-mlps-pytorch
+Provides-Extra: examples
+Provides-Extra: test
+Requires-Dist: accelerate; extra == 'test'
+Requires-Dist: diffusers>=0.32.0; extra == 'test'
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: transformers; 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
+
+## Appreciation
+
+- [Pranoy](https://github.com/pranoyr) for submitting a pull request for proprioception masking
+
+## Install
+
+```shell
+$ pip install mimic-video
+```
+
+## Usage
+
+```python
+import torch
+
+# video wrapper
+# but will be agnostic to the model
+
+from mimic_video.cosmos_predict import CosmosPredictWrapper
+
+video_wrapper = CosmosPredictWrapper(
+    extract_layer = 1,
+    random_weights = True,
+    tiny = True
+)
+
+# mimic video
+
+from mimic_video import MimicVideo
+
+model = MimicVideo(512, video_wrapper)
+
+# states
+
+video = torch.rand(2, 3, 3, 32, 32)
+
+joint_state = torch.randn(2, 32)
+
+# action
+
+actions = torch.randn(2, 32, 20)
+
+# training
+
+loss = model(
+    prompts = [
+        'put the package on the conveyer belt',
+        'pass the butter'
+    ],
+    video = video,
+    actions = actions,
+    joint_state = joint_state
+)
+
+loss.backward()
+
+# inference
+
+actions = model.sample(
+    prompts = 'peel the orange',
+    video = video[:1],
+    joint_state = joint_state[:1]
+)
+
+assert actions.shape == (1, 32, 20)
+```
+
+## 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}
+}
+```
+
+```bibtex
+@misc{black2025trainingtimeactionconditioningefficient,
+    title   = {Training-Time Action Conditioning for Efficient Real-Time Chunking}, 
+    author  = {Kevin Black and Allen Z. Ren and Michael Equi and Sergey Levine},
+    year    = {2025},
+    eprint  = {2512.05964},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2512.05964}, 
+}
+```
+
+```bibtex
+@misc{intelligence2025pi06vlalearnsexperience,
+    title   = {$\pi^{*}_{0.6}$: a VLA That Learns From Experience}, 
+    author  = {Physical Intelligence and Ali Amin and Raichelle Aniceto and Ashwin Balakrishna and Kevin Black and Ken Conley and Grace Connors and James Darpinian and Karan Dhabalia and Jared DiCarlo and Danny Driess and Michael Equi and Adnan Esmail and Yunhao Fang and Chelsea Finn and Catherine Glossop and Thomas Godden and Ivan Goryachev and Lachy Groom and Hunter Hancock and Karol Hausman and Gashon Hussein and Brian Ichter and Szymon Jakubczak and Rowan Jen and Tim Jones and Ben Katz and Liyiming Ke and Chandra Kuchi and Marinda Lamb and Devin LeBlanc and Sergey Levine and Adrian Li-Bell and Yao Lu and Vishnu Mano and Mohith Mothukuri and Suraj Nair and Karl Pertsch and Allen Z. Ren and Charvi Sharma and Lucy Xiaoyang Shi and Laura Smith and Jost Tobias Springenberg and Kyle Stachowicz and Will Stoeckle and Alex Swerdlow and James Tanner and Marcel Torne and Quan Vuong and Anna Walling and Haohuan Wang and Blake Williams and Sukwon Yoo and Lili Yu and Ury Zhilinsky and Zhiyuan Zhou},
+    year    = {2025},
+    eprint  = {2511.14759},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2511.14759}, 
+}
+```

mimic_video-0.0.31.dist-info/RECORD

@@ -0,0 +1,7 @@
+mimic_video/__init__.py,sha256=Rs3QeBBGBKKi1U1ykcyeBrCL2XCbfNvppeeD1Fb1pdY,47
+mimic_video/cosmos_predict.py,sha256=TJYIBFY2OGMAb1-5U5AOYZPbUKx-CuBIN3qbnpVWj4k,9078
+mimic_video/mimic_video.py,sha256=9jMkWzHu9qxVN5Uez8jTgErsWyaiJz9r-FVhPpf7Fbg,23165
+mimic_video-0.0.31.dist-info/METADATA,sha256=icex90WVH5NJhaZVUJ-2NI3X3qDHZhmKBxxsoe20kTQ,5876
+mimic_video-0.0.31.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mimic_video-0.0.31.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+mimic_video-0.0.31.dist-info/RECORD,,

mimic_video-0.0.31.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.31.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.