| 10 | |
| 11 | |
| 12 | class MotionCtrl(LatentDiffusion): |
| 13 | def __init__(self, |
| 14 | omcm_config=None, |
| 15 | pose_dim=12, |
| 16 | context_dim=1024, |
| 17 | *args, |
| 18 | **kwargs): |
| 19 | super(MotionCtrl, self).__init__(*args, **kwargs) |
| 20 | |
| 21 | # object motion control module |
| 22 | if omcm_config is not None: |
| 23 | self.omcm = instantiate_from_config(omcm_config) |
| 24 | else: |
| 25 | self.omcm = None |
| 26 | |
| 27 | |
| 28 | # camera motion control module |
| 29 | |
| 30 | bound_method = selfattn_forward_unet.__get__( |
| 31 | self.model.diffusion_model, |
| 32 | self.model.diffusion_model.__class__) |
| 33 | setattr(self.model.diffusion_model, 'forward', bound_method) |
| 34 | |
| 35 | for _name, _module in self.model.diffusion_model.named_modules(): |
| 36 | if _module.__class__.__name__ == 'TemporalTransformer': |
| 37 | bound_method = TemporalTransformer_forward.__get__( |
| 38 | _module, _module.__class__) |
| 39 | setattr(_module, 'forward', bound_method) |
| 40 | |
| 41 | if _module.__class__.__name__ == 'BasicTransformerBlock': |
| 42 | # SpatialTransformer only |
| 43 | if _module.attn2.to_k.in_features != context_dim: # TemporalTransformer without crossattn |
| 44 | |
| 45 | bound_method = temporal_selfattn_forward_BasicTransformerBlock.__get__( |
| 46 | _module, _module.__class__) |
| 47 | setattr(_module, '_forward', bound_method) |
| 48 | |
| 49 | cc_projection = nn.Linear(_module.attn2.to_k.in_features + pose_dim, _module.attn2.to_k.in_features) |
| 50 | nn.init.eye_(list(cc_projection.parameters())[0][:_module.attn2.to_k.in_features, :_module.attn2.to_k.in_features]) |
| 51 | nn.init.zeros_(list(cc_projection.parameters())[1]) |
| 52 | cc_projection.requires_grad_(True) |
| 53 | |
| 54 | _module.add_module('cc_projection', cc_projection) |
| 55 | |
| 56 | else: |
| 57 | bound_method = spatial_forward_BasicTransformerBlock.__get__( |
| 58 | _module, _module.__class__) |
| 59 | setattr(_module, '_forward', bound_method) |
| 60 | |
| 61 | def get_traj_features(self, extra_cond): |
| 62 | b, c, t, h, w = extra_cond.shape |
| 63 | ## process in 2D manner |
| 64 | extra_cond = rearrange(extra_cond, 'b c t h w -> (b t) c h w') |
| 65 | traj_features = self.omcm(extra_cond) |
| 66 | traj_features = [rearrange(feature, '(b t) c h w -> b c t h w', b=b, t=t) for feature in traj_features] |
| 67 | return traj_features |
nothing calls this directly
no outgoing calls
no test coverage detected