mideway

algorithms/worldmem/df_video.py

@@ -1,3 +1,4 @@
+import os
 import random
 import math
 import numpy as np
@@ -19,7 +20,7 @@ from .df_base import DiffusionForcingBase
 from .models.vae import VAE_models
 from .models.diffusion import Diffusion
 from .models.pose_prediction import PosePredictionNet
-
+import glob
 
 # Utility Functions
 def euler_to_rotation_matrix(pitch, yaw):
@@ -337,27 +338,28 @@ class WorldMemMinecraft(DiffusionForcingBase):
         self.n_frames = cfg.n_frames
         if hasattr(cfg, "n_tokens"):
             self.n_tokens = cfg.n_tokens // cfg.frame_stack
-        self.condition_similar_length = cfg.condition_similar_length
-        self.pose_cond_dim = cfg.pose_cond_dim
+        self.memory_condition_length = cfg.memory_condition_length
+        self.pose_cond_dim = getattr(cfg, "pose_cond_dim", 5)
 
         self.use_plucker = cfg.use_plucker
         self.relative_embedding = cfg.relative_embedding
-        self.cond_only_on_qk = cfg.cond_only_on_qk
-        self.use_reference_attention = cfg.use_reference_attention
-        self.add_frame_timestep_embedder = cfg.add_frame_timestep_embedder
+        self.state_embed_only_on_qk = getattr(cfg, "state_embed_only_on_qk", True)
+        self.use_memory_attention = getattr(cfg, "use_memory_attention", True)
+        self.add_timestamp_embedding = cfg.add_timestamp_embedding
         self.ref_mode = getattr(cfg, "ref_mode", 'sequential')
         self.log_curve = getattr(cfg, "log_curve", False)
-        self.focal_length = cfg.focal_length
+        self.focal_length =  getattr(cfg, "focal_length", 0.35)
         self.log_video = cfg.log_video
         self.self_consistency_eval = getattr(cfg, "self_consistency_eval", False)
-        self.next_frame_length = cfg.next_frame_length
+        self.next_frame_length = getattr(cfg, "next_frame_length", 1)
+        self.require_pose_prediction = getattr(cfg, "require_pose_prediction", False)
 
         super().__init__(cfg)
             
     def _build_model(self):
 
         self.diffusion_model = Diffusion(
-            reference_length=self.condition_similar_length,
+            reference_length=self.memory_condition_length,
             x_shape=self.x_stacked_shape,
             action_cond_dim=self.action_cond_dim,
             pose_cond_dim=self.pose_cond_dim,
@@ -366,19 +368,18 @@ class WorldMemMinecraft(DiffusionForcingBase):
             is_dit=True,
             use_plucker=self.use_plucker,
             relative_embedding=self.relative_embedding,
-            cond_only_on_qk=self.cond_only_on_qk,
-            use_reference_attention=self.use_reference_attention,
-            add_frame_timestep_embedder=self.add_frame_timestep_embedder,
+            state_embed_only_on_qk=self.state_embed_only_on_qk,
+            use_memory_attention=self.use_memory_attention,
+            add_timestamp_embedding=self.add_timestamp_embedding,
             ref_mode=self.ref_mode
         )
 
-        # self.register_data_mean_std(self.cfg.data_mean, self.cfg.data_std)
         self.validation_lpips_model = LearnedPerceptualImagePatchSimilarity()
-
         vae = VAE_models["vit-l-20-shallow-encoder"]()
         self.vae = vae.eval()
 
-        self.pose_prediction_model = PosePredictionNet()
+        if self.require_pose_prediction:
+            self.pose_prediction_model = PosePredictionNet()
 
     def _generate_noise_levels(self, xs: torch.Tensor, masks = None) -> torch.Tensor:
         """
@@ -422,7 +423,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
                 for i in range(self.n_frames):
                     input_pose_condition.append(
                         convert_to_plucker(
-                            torch.cat([c2w_mat[i:i + 1], c2w_mat[-self.condition_similar_length:]]).clone(),
+                            torch.cat([c2w_mat[i:i + 1], c2w_mat[-self.memory_condition_length:]]).clone(),
                             0,
                             focal_length=self.focal_length,
                             image_height=xs.shape[-2],image_width=xs.shape[-1]
@@ -431,7 +432,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
                     frame_idx_list.append(
                         torch.cat([
                             frame_idx[i:i + 1] - frame_idx[i:i + 1],
-                            frame_idx[-self.condition_similar_length:] - frame_idx[i:i + 1]
+                            frame_idx[-self.memory_condition_length:] - frame_idx[i:i + 1]
                         ]).clone()
                     )
                 input_pose_condition = torch.cat(input_pose_condition)
@@ -449,21 +450,21 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
         noise_levels = self._generate_noise_levels(xs)
 
-        if self.condition_similar_length:
-            noise_levels[-self.condition_similar_length:] = self.diffusion_model.stabilization_level
-            conditions[-self.condition_similar_length:] *= 0
+        if self.memory_condition_length:
+            noise_levels[-self.memory_condition_length:] = self.diffusion_model.stabilization_level
+            conditions[-self.memory_condition_length:] *= 0
 
         _, loss = self.diffusion_model(
             xs,
             conditions,
             input_pose_condition,
             noise_levels=noise_levels,
-            reference_length=self.condition_similar_length,
+            reference_length=self.memory_condition_length,
             frame_idx=frame_idx_list
         )
 
-        if self.condition_similar_length:
-            loss = loss[:-self.condition_similar_length]
+        if self.memory_condition_length:
+            loss = loss[:-self.memory_condition_length]
 
         loss = self.reweight_loss(loss, None)
 
@@ -472,7 +473,6 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
         return {"loss": loss}
     
-
     def on_validation_epoch_end(self, namespace="validation") -> None:
         if not self.validation_step_outputs:
             return
@@ -577,12 +577,12 @@ class WorldMemMinecraft(DiffusionForcingBase):
         x = rearrange(x, "(t b) c h w-> t b c h w", t=total_frames)
         return x
 
-    def _generate_condition_indices(self, curr_frame, condition_similar_length, xs_pred, pose_conditions, frame_idx, horizon):
+    def _generate_condition_indices(self, curr_frame, memory_condition_length, xs_pred, pose_conditions, frame_idx, horizon):
         """
         Generate indices for condition similarity based on the current frame and pose conditions.
         """
-        if curr_frame < condition_similar_length:
-            random_idx = [i for i in range(curr_frame)] + [0] * (condition_similar_length - curr_frame)
+        if curr_frame < memory_condition_length:
+            random_idx = [i for i in range(curr_frame)] + [0] * (memory_condition_length - curr_frame)
             random_idx = np.repeat(np.array(random_idx)[:, None], xs_pred.shape[1], -1)
         else:
             # Generate points in a sphere and filter based on field of view
@@ -614,7 +614,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
             ])
 
             random_idx = []
-            for _ in range(condition_similar_length):
+            for _ in range(memory_condition_length):
                 overlap_ratio = ((in_fov1.bool() & in_fov_list).sum(1)) / in_fov1.sum()
                 
                 confidence = overlap_ratio + (curr_frame - frame_idx[:curr_frame]) / curr_frame * (-0.2)
@@ -674,15 +674,15 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
         return input_condition, input_pose_condition, frame_idx_list
 
-    def _prepare_noise_levels(self, scheduling_matrix, m, curr_frame, batch_size, condition_similar_length):
+    def _prepare_noise_levels(self, scheduling_matrix, m, curr_frame, batch_size, memory_condition_length):
         """
         Prepare noise levels for the current sampling step.
         """
         from_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m]))[:, None].repeat(batch_size, axis=1)
         to_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m + 1]))[:, None].repeat(batch_size, axis=1)
-        if condition_similar_length:
-            from_noise_levels = np.concatenate([from_noise_levels, np.zeros((condition_similar_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
-            to_noise_levels = np.concatenate([to_noise_levels, np.zeros((condition_similar_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
+        if memory_condition_length:
+            from_noise_levels = np.concatenate([from_noise_levels, np.zeros((memory_condition_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
+            to_noise_levels = np.concatenate([to_noise_levels, np.zeros((memory_condition_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
         from_noise_levels = torch.from_numpy(from_noise_levels).to(self.device)
         to_noise_levels = torch.from_numpy(to_noise_levels).to(self.device)
         return from_noise_levels, to_noise_levels
@@ -703,7 +703,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
             None: Appends the predicted and ground truth frames to `self.validation_step_outputs`.
         """
         # Preprocess the input batch
-        condition_similar_length = self.condition_similar_length
+        memory_condition_length = self.memory_condition_length
         xs_raw, conditions, pose_conditions, c2w_mat, frame_idx = self._preprocess_batch(batch)
 
 
@@ -744,9 +744,9 @@ class WorldMemMinecraft(DiffusionForcingBase):
             pbar.set_postfix({"start": start_frame, "end": curr_frame + horizon})
 
             # Handle condition similarity logic
-            if condition_similar_length:
+            if memory_condition_length:
                 random_idx = self._generate_condition_indices(
-                    curr_frame, condition_similar_length, xs_pred, pose_conditions, frame_idx
+                    curr_frame, memory_condition_length, xs_pred, pose_conditions, frame_idx, horizon
                 )
 
                 xs_pred = torch.cat([xs_pred, xs_pred[random_idx[:, range(xs_pred.shape[1])], range(xs_pred.shape[1])].clone()], 0)
@@ -760,7 +760,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
             # Perform sampling for each step in the scheduling matrix
             for m in range(scheduling_matrix.shape[0] - 1):
                 from_noise_levels, to_noise_levels = self._prepare_noise_levels(
-                    scheduling_matrix, m, curr_frame, batch_size, condition_similar_length
+                    scheduling_matrix, m, curr_frame, batch_size, memory_condition_length
                 )
 
                 xs_pred[start_frame:] = self.diffusion_model.sample_step(
@@ -771,13 +771,13 @@ class WorldMemMinecraft(DiffusionForcingBase):
                     to_noise_levels[start_frame:],
                     current_frame=curr_frame,
                     mode="validation",
-                    reference_length=condition_similar_length,
+                    reference_length=memory_condition_length,
                     frame_idx=frame_idx_list
                 ).cpu()
 
             # Remove condition similarity frames if applicable
-            if condition_similar_length:
-                xs_pred = xs_pred[:-condition_similar_length]
+            if memory_condition_length:
+                xs_pred = xs_pred[:-memory_condition_length]
 
             curr_frame += horizon
             pbar.update(horizon)
@@ -794,7 +794,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
     def interactive(self, first_frame, new_actions, first_pose, device,
                     memory_latent_frames, memory_actions, memory_poses, memory_c2w, memory_frame_idx):
     
-        condition_similar_length = self.condition_similar_length
+        memory_condition_length = self.memory_condition_length
 
         if memory_latent_frames is None:
             first_frame = torch.from_numpy(first_frame)
@@ -881,13 +881,13 @@ class WorldMemMinecraft(DiffusionForcingBase):
             )
 
             # Handle condition similarity logic
-            if condition_similar_length:
+            if memory_condition_length:
                 random_idx = self._generate_condition_indices(
-                    curr_frame, condition_similar_length, xs_pred, pose_conditions, frame_idx, next_horizon
+                    curr_frame, memory_condition_length, xs_pred, pose_conditions, frame_idx, next_horizon
                 )
                 
                 # random_idx = np.unique(random_idx)[:, None]
-                # condition_similar_length = len(random_idx)
+                # memory_condition_length = len(random_idx)
                 xs_pred = torch.cat([xs_pred, xs_pred[random_idx[:, range(xs_pred.shape[1])], range(xs_pred.shape[1])].clone()], 0)
 
             # Prepare input conditions and pose conditions
@@ -899,7 +899,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
             # Perform sampling for each step in the scheduling matrix
             for m in range(scheduling_matrix.shape[0] - 1):
                 from_noise_levels, to_noise_levels = self._prepare_noise_levels(
-                    scheduling_matrix, m, curr_frame, batch_size, condition_similar_length
+                    scheduling_matrix, m, curr_frame, batch_size, memory_condition_length
                 )
 
                 xs_pred[start_frame:] = self.diffusion_model.sample_step(
@@ -910,13 +910,13 @@ class WorldMemMinecraft(DiffusionForcingBase):
                     to_noise_levels[start_frame:],
                     current_frame=curr_frame,
                     mode="validation",
-                    reference_length=condition_similar_length,
+                    reference_length=memory_condition_length,
                     frame_idx=frame_idx_list
                 ).cpu()
 
 
-            if condition_similar_length:
-                xs_pred = xs_pred[:-condition_similar_length]
+            if memory_condition_length:
+                xs_pred = xs_pred[:-memory_condition_length]
 
             curr_frame += next_horizon
             pbar.update(next_horizon)
@@ -926,4 +926,4 @@ class WorldMemMinecraft(DiffusionForcingBase):
         xs_pred = self.decode(xs_pred[n_context_frames:].to(device)).cpu()
 
         return xs_pred.cpu().numpy(), memory_latent_frames.cpu().numpy(), memory_actions.cpu().numpy(), \
-            memory_poses.cpu().numpy(), memory_c2w.cpu().numpy(), memory_frame_idx.cpu().numpy()
+            memory_poses.cpu().numpy(), memory_c2w.cpu().numpy(), memory_frame_idx.cpu().numpy()

algorithms/worldmem/models/attention.py

@@ -288,12 +288,12 @@ class MemFullAttention(nn.Module):
 
     def forward(self, x: torch.Tensor, relative_embedding=False,
                 extra_condition=None,
-                cond_only_on_qk=False,
+                state_embed_only_on_qk=False,
                 reference_length=None):
 
         B, T, H, W, D = x.shape
 
-        if cond_only_on_qk:
+        if state_embed_only_on_qk:
             q, k, _ = self.to_qkv(x+extra_condition).chunk(3, dim=-1)
             _, _, v = self.to_qkv(x).chunk(3, dim=-1)
         else:

algorithms/worldmem/models/diffusion.py

@@ -26,9 +26,9 @@ class Diffusion(nn.Module):
         is_dit: bool=False,
         use_plucker=False,
         relative_embedding=False,
-        cond_only_on_qk=False,
-        use_reference_attention=False,
-        add_frame_timestep_embedder=False,
+        state_embed_only_on_qk=False,
+        use_memory_attention=False,
+        add_timestamp_embedding=False,
         ref_mode='sequential'
     ):
         super().__init__()
@@ -54,9 +54,9 @@ class Diffusion(nn.Module):
         self.pose_cond_dim = pose_cond_dim
         self.use_plucker = use_plucker
         self.relative_embedding = relative_embedding
-        self.cond_only_on_qk = cond_only_on_qk
-        self.use_reference_attention = use_reference_attention
-        self.add_frame_timestep_embedder = add_frame_timestep_embedder
+        self.state_embed_only_on_qk = state_embed_only_on_qk
+        self.use_memory_attention = use_memory_attention
+        self.add_timestamp_embedding = add_timestamp_embedding
         self.ref_mode = ref_mode
 
         self._build_model()
@@ -69,9 +69,9 @@ class Diffusion(nn.Module):
                                             pose_cond_dim=self.pose_cond_dim, reference_length=self.reference_length,
                                             use_plucker=self.use_plucker,
                                             relative_embedding=self.relative_embedding,
-                                            cond_only_on_qk=self.cond_only_on_qk,
-                                            use_reference_attention=self.use_reference_attention,
-                                            add_frame_timestep_embedder=self.add_frame_timestep_embedder,
+                                            state_embed_only_on_qk=self.state_embed_only_on_qk,
+                                            use_memory_attention=self.use_memory_attention,
+                                            add_timestamp_embedding=self.add_timestamp_embedding,
                                             ref_mode=self.ref_mode)
         else:
             raise NotImplementedError

algorithms/worldmem/models/dit.py

@@ -153,8 +153,8 @@ class SpatioTemporalDiTBlock(nn.Module):
         reference_rotary_emb=None,
         use_plucker=False,
         relative_embedding=False,
-        cond_only_on_qk=False,
-        use_reference_attention=False,
+        state_embed_only_on_qk=False,
+        use_memory_attention=False,
         ref_mode='sequential'
     ):
         super().__init__()
@@ -196,8 +196,8 @@ class SpatioTemporalDiTBlock(nn.Module):
         )
         self.t_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
 
-        self.use_reference_attention = use_reference_attention
-        if self.use_reference_attention:
+        self.use_memory_attention = use_memory_attention
+        if self.use_memory_attention:
             self.r_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
             self.ref_type = "full_ref"
             if self.ref_type == "temporal_ref":
@@ -233,7 +233,7 @@ class SpatioTemporalDiTBlock(nn.Module):
 
         self.reference_length = reference_length
         self.relative_embedding = relative_embedding
-        self.cond_only_on_qk = cond_only_on_qk
+        self.state_embed_only_on_qk = state_embed_only_on_qk
 
         self.ref_mode = ref_mode
 
@@ -265,7 +265,7 @@ class SpatioTemporalDiTBlock(nn.Module):
         # memory block
         relative_embedding = self.relative_embedding # and mode == "training"
 
-        if self.use_reference_attention:
+        if self.use_memory_attention:
             r_shift_msa, r_scale_msa, r_gate_msa, r_shift_mlp, r_scale_mlp, r_gate_mlp = self.r_adaLN_modulation(c).chunk(6, dim=-1)
 
             if pose_cond is not None:
@@ -288,7 +288,7 @@ class SpatioTemporalDiTBlock(nn.Module):
                         # if current_frame == 18:
                         #     import pdb;pdb.set_trace()
 
-                        if self.cond_only_on_qk:
+                        if self.state_embed_only_on_qk:
                             attn_input = x1_zero_frame
                             extra_condition = input_cond                        
                         else:
@@ -304,7 +304,7 @@ class SpatioTemporalDiTBlock(nn.Module):
                     x = x + gate(self.r_attn(modulate(self.r_norm1(attn_input), r_shift_msa, r_scale_msa), 
                                     relative_embedding=relative_embedding, 
                                     extra_condition=extra_condition, 
-                                    cond_only_on_qk=self.cond_only_on_qk,
+                                    state_embed_only_on_qk=self.state_embed_only_on_qk,
                                     reference_length=reference_length), r_gate_msa)
                 else:
                     # pose_cond *= 0
@@ -354,9 +354,9 @@ class DiT(nn.Module):
         reference_length=8,
         use_plucker=False,
         relative_embedding=False,
-        cond_only_on_qk=False,
-        use_reference_attention=False,
-        add_frame_timestep_embedder=False,
+        state_embed_only_on_qk=False,
+        use_memory_attention=False,
+        add_timestamp_embedding=False,
         ref_mode='sequential'
     ):
         super().__init__()
@@ -369,9 +369,9 @@ class DiT(nn.Module):
         self.x_embedder = PatchEmbed(input_h, input_w, patch_size, in_channels, hidden_size, flatten=False)
         self.t_embedder = TimestepEmbedder(hidden_size)
 
-        self.add_frame_timestep_embedder = add_frame_timestep_embedder
-        if self.add_frame_timestep_embedder:
-            self.frame_timestep_embedder = TimestepEmbedder(hidden_size)
+        self.add_timestamp_embedding = add_timestamp_embedding
+        if self.add_timestamp_embedding:
+            self.timestamp_embedding = TimestepEmbedder(hidden_size)
 
         frame_h, frame_w = self.x_embedder.grid_size
 
@@ -404,14 +404,14 @@ class DiT(nn.Module):
                     reference_rotary_emb=self.reference_rotary_emb,
                     use_plucker=self.use_plucker,
                     relative_embedding=relative_embedding,
-                    cond_only_on_qk=cond_only_on_qk,
-                    use_reference_attention=use_reference_attention,
+                    state_embed_only_on_qk=state_embed_only_on_qk,
+                    use_memory_attention=use_memory_attention,
                     ref_mode=ref_mode
                 )
                 for _ in range(depth)
             ]
         )
-        self.use_reference_attention = use_reference_attention
+        self.use_memory_attention = use_memory_attention
         self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
         self.initialize_weights()
 
@@ -434,7 +434,7 @@ class DiT(nn.Module):
         nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
         nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
 
-        if self.use_reference_attention:
+        if self.use_memory_attention:
             if not self.use_plucker:
                 nn.init.normal_(self.position_embedder.mlp[0].weight, std=0.02)
                 nn.init.normal_(self.position_embedder.mlp[2].weight, std=0.02)
@@ -442,9 +442,9 @@ class DiT(nn.Module):
                 nn.init.normal_(self.angle_embedder.mlp[0].weight, std=0.02)
                 nn.init.normal_(self.angle_embedder.mlp[2].weight, std=0.02)
             
-            if self.add_frame_timestep_embedder:
-                nn.init.normal_(self.frame_timestep_embedder.mlp[0].weight, std=0.02)
-                nn.init.normal_(self.frame_timestep_embedder.mlp[2].weight, std=0.02)
+            if self.add_timestamp_embedding:
+                nn.init.normal_(self.timestamp_embedding.mlp[0].weight, std=0.02)
+                nn.init.normal_(self.timestamp_embedding.mlp[2].weight, std=0.02)
 
 
         # Zero-out adaLN modulation layers in DiT blocks:
@@ -454,7 +454,7 @@ class DiT(nn.Module):
             nn.init.constant_(block.t_adaLN_modulation[-1].weight, 0)
             nn.init.constant_(block.t_adaLN_modulation[-1].bias, 0)
 
-            if self.use_plucker and self.use_reference_attention:
+            if self.use_plucker and self.use_memory_attention:
                 nn.init.constant_(block.pose_cond_mlp.weight, 0)
                 nn.init.constant_(block.pose_cond_mlp.bias, 0)
 
@@ -526,10 +526,10 @@ class DiT(nn.Module):
                 pc = self.pose_embedder(pose_cond)
                 pc = pc.permute(1,0,2,3,4)
 
-                if torch.is_tensor(frame_idx) and self.add_frame_timestep_embedder:
+                if torch.is_tensor(frame_idx) and self.add_timestamp_embedding:
                     bb = frame_idx.shape[1]
                     frame_idx = rearrange(frame_idx, "t b -> (b t)")
-                    frame_idx = self.frame_timestep_embedder(frame_idx)
+                    frame_idx = self.timestamp_embedding(frame_idx)
                     frame_idx = rearrange(frame_idx, "(b t) d -> b t d", b=bb)
                     pc = pc + frame_idx[:, :, None, None]
 
@@ -545,17 +545,12 @@ class DiT(nn.Module):
         x = rearrange(x, "b t h w d -> (b t) h w d")
         x = self.unpatchify(x)  # (N, out_channels, H, W)
         x = rearrange(x, "(b t) c h w -> b t c h w", t=T)
-
-        # print("self.blocks[0].pose_cond_mlp.weight:", self.blocks[0].pose_cond_mlp.weight)
-        # print("self.blocks[0].r_adaLN_modulation[1].weight:", self.blocks[0].r_adaLN_modulation[1].weight)
-        # print("self.blocks[0].t_adaLN_modulation[1].weight:", self.blocks[0].t_adaLN_modulation[1].weight)
-
         return x
 
 
 def DiT_S_2(action_cond_dim, pose_cond_dim, reference_length, 
 use_plucker, relative_embedding, 
-cond_only_on_qk, use_reference_attention, add_frame_timestep_embedder,
+state_embed_only_on_qk, use_memory_attention, add_timestamp_embedding,
 ref_mode):
     return DiT(
         patch_size=2,
@@ -567,9 +562,9 @@ ref_mode):
         reference_length=reference_length,
         use_plucker=use_plucker,
         relative_embedding=relative_embedding,
-        cond_only_on_qk=cond_only_on_qk,
-        use_reference_attention=use_reference_attention,
-        add_frame_timestep_embedder=add_frame_timestep_embedder,
+        state_embed_only_on_qk=state_embed_only_on_qk,
+        use_memory_attention=use_memory_attention,
+        add_timestamp_embedding=add_timestamp_embedding,
         ref_mode=ref_mode
     )
 

app.py

@@ -34,11 +34,23 @@ def load_custom_checkpoint(algo, checkpoint_path):
         model_path = hf_hub_download(repo_id=repo_id, 
                             filename=file_name)
         ckpt = torch.load(model_path, map_location=torch.device('cpu'))
-        algo.load_state_dict(ckpt['state_dict'], strict=False)
+
+        algo.load_state_dict(ckpt['state_dict'], strict=True)
         print("Load: ", model_path)
     except:
         ckpt = torch.load(checkpoint_path, map_location=torch.device('cpu'))
-        algo.load_state_dict(ckpt['state_dict'], strict=False)      
+
+        filtered_state_dict = {}
+        for k, v in ckpt['state_dict'].items():
+            if "frame_timestep_embedder" in k:
+                new_k = k.replace("frame_timestep_embedder", "timestamp_embedding")
+                filtered_state_dict[new_k] = v
+            else:
+                filtered_state_dict[k] = v
+
+        algo.load_state_dict(filtered_state_dict, strict=True)
+
+        # algo.load_state_dict(ckpt['state_dict'], strict=True)      
         print("Load: ", checkpoint_path)  
 
 def download_assets_if_needed():
@@ -233,11 +245,11 @@ def set_context_length(context_length, sampling_context_length_state):
     print("set context length to", worldmem.n_tokens)
     return sampling_context_length_state
 
-def set_memory_length(memory_length, sampling_memory_length_state):
-    worldmem.condition_similar_length = memory_length
-    sampling_memory_length_state = memory_length
-    print("set memory length to", worldmem.condition_similar_length)
-    return sampling_memory_length_state
+def set_memory_condition_length(memory_condition_length, sampling_memory_condition_length_state):
+    worldmem.memory_condition_length = memory_condition_length
+    sampling_memory_condition_length_state = memory_condition_length
+    print("set memory length to", worldmem.memory_condition_length)
+    return sampling_memory_condition_length_state
 
 def set_next_frame_length(next_frame_length, sampling_next_frame_length_state):
     worldmem.next_frame_length = next_frame_length
@@ -509,8 +521,8 @@ with gr.Blocks(css=css) as demo:
                 label="Context Length",
                 info="How many previous frames in temporal context window."
             )
-            slider_memory_length = gr.Slider(
-                minimum=4, maximum=16, value=worldmem.condition_similar_length, step=1,
+            slider_memory_condition_length = gr.Slider(
+                minimum=4, maximum=16, value=worldmem.memory_condition_length, step=1,
                 label="Memory Length",
                 info="How many previous frames in memory window. (Recommended: 1, multi-frame generation is not stable yet)"
             )
@@ -522,7 +534,7 @@ with gr.Blocks(css=css) as demo:
     
     sampling_timesteps_state = gr.State(worldmem.sampling_timesteps)
     sampling_context_length_state = gr.State(worldmem.n_tokens)
-    sampling_memory_length_state = gr.State(worldmem.condition_similar_length)
+    sampling_memory_condition_length_state = gr.State(worldmem.memory_condition_length)
     sampling_next_frame_length_state = gr.State(worldmem.next_frame_length)
 
     video_frames = gr.State(load_image_as_tensor(selected_image.value)[None].numpy())
@@ -547,7 +559,7 @@ with gr.Blocks(css=css) as demo:
 
     examples = gr.Examples(
         examples=example_images,
-        inputs=[example_case, image_output, log_output, slider_denoising_step, slider_context_length, slider_memory_length],
+        inputs=[example_case, image_output, log_output, slider_denoising_step, slider_context_length, slider_memory_condition_length],
         cache_examples=False
     )
 
@@ -574,7 +586,7 @@ with gr.Blocks(css=css) as demo:
 
     slider_denoising_step.change(fn=set_denoising_steps, inputs=[slider_denoising_step, sampling_timesteps_state], outputs=sampling_timesteps_state)
     slider_context_length.change(fn=set_context_length, inputs=[slider_context_length, sampling_context_length_state], outputs=sampling_context_length_state)
-    slider_memory_length.change(fn=set_memory_length, inputs=[slider_memory_length, sampling_memory_length_state], outputs=sampling_memory_length_state)
+    slider_memory_condition_length.change(fn=set_memory_condition_length, inputs=[slider_memory_condition_length, sampling_memory_condition_length_state], outputs=sampling_memory_condition_length_state)
     slider_next_frame_length.change(fn=set_next_frame_length, inputs=[slider_next_frame_length, sampling_next_frame_length_state], outputs=sampling_next_frame_length_state)
 
-demo.launch()
+demo.launch(share=True)

configurations/huggingface.yaml

@@ -3,14 +3,15 @@ pose_cond_dim: 5
 use_plucker: true
 focal_length: 0.35
 customized_validation: true
-condition_similar_length: 8
+memory_condition_length: 8
 log_video: true
 relative_embedding: true
-cond_only_on_qk: true
-add_pose_embed: false
+state_embed_only_on_qk: true
 use_domain_adapter: false
-use_reference_attention: true
-add_frame_timestep_embedder: true
+use_memory_attention: true
+add_timestamp_embedding: true
+use_pose_prediction: true
+require_pose_prediction: true
 is_interactive: true
 diffusion:
   sampling_timesteps: 20