mimic-video 0.0.1__py3-none-any.whl → 0.0.19__py3-none-any.whl

mimic_video/__init__.py

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

mimic_video/cosmos_predict.py

@@ -0,0 +1,269 @@
+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_layer: int = 19,
+        random_weights: bool = False,
+        tiny: bool = False,
+        normalize = lambda t: (t - 0.5) * 2.0
+    ):
+        super().__init__()
+        self.extract_layer = extract_layer
+        self.hook_handle = None
+        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, self.extract_layer + 1)
+        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 exists(self.hook_handle):
+            self.hook_handle.remove()
+
+    def _register_hook(self):
+        assert hasattr(self.transformer, 'transformer_blocks'), 'transformer must have transformer_blocks'
+        assert len(self.transformer.transformer_blocks) > self.extract_layer, f'layer {self.extract_layer} out of bounds'
+
+        target_layer = self.transformer.transformer_blocks[self.extract_layer]
+
+        def hook_fn(module, inp, out):
+            self.cached_hidden_states.append(out.detach().cpu())
+
+        self.hook_handle = target_layer.register_forward_hook(hook_fn)
+
+    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) > 0, 'hidden states not captured'
+        
+        # Return hidden states from the first denoising step
+        hidden = self.cached_hidden_states[0]
+        
+        assert hidden.shape[-1] == self.dim_latent, f'hidden dim mismatch: expected {self.dim_latent_hidden}, got {hidden.shape[-1]}'
+        
+        return hidden

mimic_video/mimic_video.py

@@ -1,18 +1,24 @@
+from __future__ import annotations
+
 import torch
-from torch import nn
-from torch.nn import Module, ModuleList, Linear
+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
+from einops import einsum, rearrange, repeat
 from einops.layers.torch import Rearrange
 
 from x_mlps_pytorch import create_mlp
 
+from tqdm import tqdm
+
 from torch_einops_utils import (
     pad_left_ndim,
-    align_dims_left
+    align_dims_left,
+    pad_at_dim,
+    pack_with_inverse,
 )
 
 # ein notation
@@ -37,12 +43,23 @@ def divisible_by(num, den):
 
 # 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)
+
 # time
 
 # they follow p0's research finding with the beta distribution
@@ -72,7 +89,8 @@ class AdaptiveRMSNorm(Module):
         self,
         dim,
         dim_time_cond,
-        eps = 1e-6
+        eps = 1e-6,
+        ada_ln_zero_bias = -5.
     ):
         super().__init__()
         self.scale = dim ** 0.5
@@ -83,6 +101,8 @@ class AdaptiveRMSNorm(Module):
 
         nn.init.zeros_(self.to_modulation.weight)
 
+        self.ada_ln_zero_bias = ada_ln_zero_bias
+
     def forward(
         self,
         tokens,
@@ -100,7 +120,9 @@ class AdaptiveRMSNorm(Module):
 
         adaptive_normed = normed * (scale + 1.) + shift
 
-        return adaptive_normed, gate
+        gate_with_bias = (gate + self.ada_ln_zero_bias).sigmoid()
+
+        return adaptive_normed, gate_with_bias
 
 # attention
 
@@ -136,15 +158,20 @@ class Attention(Module):
         self,
         tokens,
         context = None,
-        context_mask = None
+        context_mask = None,
+        kv = None,
+        return_kv = False
     ):
         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))
+
+        if not exists(kv):
+            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
 
@@ -160,7 +187,12 @@ class Attention(Module):
 
         out = self.merge_heads(out)
 
-        return self.to_out(out)
+        out = self.to_out(out)
+
+        if not return_kv:
+            return out
+
+        return out, stack((keys, values))
 
 # feedforward
 
@@ -193,18 +225,43 @@ class MimicVideo(Module):
     def __init__(
         self,
         dim,
+        video_predict_wrapper: Module | None = None,
         *,
-        dim_video_hidden,
+        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
     ):
         super().__init__()
 
+        self.depth = depth
+
+        # maybe video predict
+
+        self.video_predict_wrapper = video_predict_wrapper
+
+        # dims
+
+        self.action_chunk_len = action_chunk_len
+        self.dim_action = dim_action
+
+        self.action_shape = (action_chunk_len, dim_action)
+        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
+
         # flow related
 
         self.sample_time_fn = default(sample_time_fn, default_sample_time_fn)
@@ -215,13 +272,26 @@ class MimicVideo(Module):
 
         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.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)
 
+        # rnn
+
+        self.rnn = GRU(dim, dim)
+
         # transformer
 
         layers = []
@@ -235,7 +305,7 @@ class MimicVideo(Module):
 
             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_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)
 
@@ -254,25 +324,93 @@ class MimicVideo(Module):
 
         self.to_pred_action_flow = nn.Sequential(
             nn.RMSNorm(dim),
-            Linear(dim, dim_action)
+            Linear(dim, dim_action, bias = False)
         )
 
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
+    @property
+    def device(self):
+        return self.zero.device
+
+    @torch.no_grad()
+    def sample(
+        self,
+        steps = 16,
+        batch_size = 1,
+        disable_progress_bar = False,
+        **kwargs
+    ):
+
+        self.eval()
+
+        noise = torch.randn((batch_size, *self.action_shape), device = self.device)
+
+        times = torch.linspace(0., 1., steps + 1, device = self.device)[:-1]
+        delta = 1. / steps
+
+        denoised = noise
+
+        cache = None
+
+        for time in tqdm(times, disable = disable_progress_bar):
+            pred_flow, cache = self.forward(actions = denoised, time = time, cache = cache, return_cache = True, **kwargs)
+
+            denoised = denoised + delta * pred_flow
+
+        return denoised
+
     def forward(
         self,
-        actions,
-        video_hiddens, # they use layer 19 of cosmos predict, at first denoising step. that's all
         *,
-        time = None,
+        actions,
+        joint_state,
+        video = None,
+        video_hiddens = None, # they use layer 19 of cosmos predict, at first denoising step. that's all
         context_mask = None,
+        time = None,
+        time_video_denoise = 0., # 0 is noise in the scheme i prefer - default to their optimal choice, but can be changed
+        prompts = None,
+        prompt_token_ids = None,
+        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
+
+        batch, device = actions.shape[0], actions.device
+
+        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_hiddens = self.video_predict_wrapper(video, prompts = prompts, prompt_token_ids = prompt_token_ids)
+                video_hiddens = video_hiddens.float() # maybe bfloat to float32
+
+                video_hiddens, _ = pack_with_inverse(video_hiddens, 'b * d')
+
+                assert video_hiddens.shape[-1] == self.dim_video_hidden
+
+            # handle video hiddens
+
+            video_hiddens = self.video_hidden_norm(video_hiddens)
 
-        is_training = not exists(time)
+        # 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:
-            batch, device = actions.shape[0], actions.device
-
             time = torch.rand((batch,), device = device)
             time = self.sample_time_fn(time)
 
@@ -285,26 +423,61 @@ class MimicVideo(Module):
         else:
             noised = actions
 
-        time_cond = self.to_time_cond(time)
+        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)
 
-        # handle video hiddens
+        times = stack((time, time_video_denoise), dim = -1)
 
-        video_hiddens = self.video_hidden_norm(video_hiddens)
+        # 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)
 
         # embed
 
         tokens = self.to_action_tokens(noised)
 
+        # one layer of rnn for actions
+
+        rnn_out, _, = self.rnn(tokens)
+        tokens = rnn_out + tokens
+
+        #  mask joint state token for proprioception masking training
+
+        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)
+
+        # pack joint with action tokens
+
+        tokens, inverse_pack = pack_with_inverse((joint_state_token, tokens), 'b * d')
+
         # transformer layers
 
-        for (
+        for ((
             attn_norm,
             attn,
             cross_attn_norm,
             cross_attn,
             ff_norm,
             ff
-        ) in self.layers:
+        ), cached_video_kv) in zip(self.layers, prev_cached_video_hiddens_kv):
 
             # cross attention
 
@@ -312,7 +485,12 @@ class MimicVideo(Module):
 
             tokens, gate = cross_attn_norm(tokens, time_cond)
 
-            tokens = residual + cross_attn(tokens, context = video_hiddens, context_mask = context_mask) * gate
+            cross_attn_out, video_kv = cross_attn(tokens, context = video_hiddens, context_mask = context_mask, kv = cached_video_kv, return_kv = True)
+
+            tokens = residual + cross_attn_out * gate
+
+            if return_cache:
+                next_cached_video_hiddens_kv.append(video_kv)
 
             # self attention
 
@@ -322,22 +500,46 @@ class MimicVideo(Module):
 
             tokens = residual + attn(tokens) * gate
 
-            # feedforward
+            # prepare feedforward
 
             residual = 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
+
+            joint_state_token, tokens = inverse_pack(tokens)
+
+            tokens = shift_feature_dim(tokens)
+
+            tokens, _ = pack_with_inverse((joint_state_token, tokens), 'b * d')
+
+            # feedforward
+
             tokens = residual + ff(tokens) * gate
 
+        # remove joint token
+
+        _, tokens = inverse_pack(tokens)
+
         # prediction
 
         pred_flow = self.to_pred_action_flow(tokens)
 
         if not is_training:
-            return pred_flow
+            # flow
+
+            out = pred_flow
+        else:
+            # mse flow loss
+
+            flow_loss = F.mse_loss(pred_flow, flow)
+
+            out = flow_loss
+
+        if not return_cache:
+            return out
 
-        # mse flow loss
+        # handle returning of cache
 
-        flow_loss = F.mse_loss(pred_flow, flow)
-        return flow_loss
+        return out, next_cached_video_hiddens_kv

{mimic_video-0.0.1.dist-info → mimic_video-0.0.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mimic-video
-Version: 0.0.1
+Version: 0.0.19
 Summary: Mimic Video
 Project-URL: Homepage, https://pypi.org/project/mimic-video/
 Project-URL: Repository, https://github.com/lucidrains/mimic-video
@@ -38,10 +38,14 @@ 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: 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>
@@ -50,6 +54,73 @@ Description-Content-Type: text/markdown
 
 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

mimic_video-0.0.19.dist-info/RECORD

@@ -0,0 +1,7 @@
+mimic_video/__init__.py,sha256=Rs3QeBBGBKKi1U1ykcyeBrCL2XCbfNvppeeD1Fb1pdY,47
+mimic_video/cosmos_predict.py,sha256=2XR9cqcUC4gKpjEDBy-GtLtMkLXvs8yKe7w8g6EeS6s,8471
+mimic_video/mimic_video.py,sha256=wQNfdV2PGlR_-S-4Mm3cyswtRQ3nBQGJiHptya3ckKU,14761
+mimic_video-0.0.19.dist-info/METADATA,sha256=67a7iVIkf557qMos_yvpcqL8dK9yL0Jf1DPQuhb2bwo,4142
+mimic_video-0.0.19.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mimic_video-0.0.19.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+mimic_video-0.0.19.dist-info/RECORD,,

mimic_video-0.0.1.dist-info/RECORD

@@ -1,6 +0,0 @@
-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,,