| 18 | |
| 19 | |
| 20 | class MLLMInContextConfig(PretrainedConfig): |
| 21 | def __init__( |
| 22 | self, |
| 23 | mllm_id: str = "Alibaba-DAMO-Academy/RynnBrain-2B", |
| 24 | diffusion_model_id: str = "SJTU-DENG-Lab/Sana_600M_512px_diffusers_64channels", |
| 25 | in_channels: int = 32, |
| 26 | input_size: int = 32, |
| 27 | num_metaqueries: int = 32, |
| 28 | _gradient_checkpointing: bool = True, |
| 29 | max_input_text_tokens: int = 256, |
| 30 | connector_num_hidden_layers: int = 12, |
| 31 | system_prompt: str = "", |
| 32 | action_condition_type: str = "no_action_condition", |
| 33 | **kwargs, |
| 34 | ): |
| 35 | super().__init__() |
| 36 | self.mllm_id = mllm_id |
| 37 | self.diffusion_model_id = diffusion_model_id |
| 38 | self.in_channels = in_channels |
| 39 | self.input_size = input_size |
| 40 | self.num_metaqueries = num_metaqueries |
| 41 | self._gradient_checkpointing = _gradient_checkpointing |
| 42 | self.max_input_text_tokens = max_input_text_tokens |
| 43 | self.connector_num_hidden_layers = connector_num_hidden_layers |
| 44 | self.system_prompt = system_prompt |
| 45 | self.action_condition_type = action_condition_type |
| 46 | |
| 47 | self.max_action_dim = kwargs.get("max_action_dim") |
| 48 | self.max_state_dim = kwargs.get("max_state_dim") |
| 49 | self.chunk_size = kwargs.get("chunk_size") |
| 50 | self.use_history_obs = kwargs.get("use_history_obs") |
| 51 | self.training_mode = kwargs.get("training_mode") |
| 52 | self.use_begin_frame_context = kwargs.get("use_begin_frame_context", False) |
| 53 | |
| 54 | |
| 55 | self.num_inference_timesteps = kwargs.get("num_inference_timesteps") |
| 56 | self.num_target_vision_tokens = kwargs.get("num_target_vision_tokens") |
| 57 | self.add_pos_embed = kwargs.get("add_pos_embed") |
| 58 | self.max_seq_len = kwargs.get("max_seq_len") |
| 59 | self.noise_beta_alpha = kwargs.get("noise_beta_alpha") |
| 60 | self.noise_beta_beta = kwargs.get("noise_beta_beta") |
| 61 | self.num_timestep_buckets = kwargs.get("num_timestep_buckets") |
| 62 | self.noise_s = kwargs.get("noise_s") |
| 63 | self.diffusion_model_cfg = kwargs.get("diffusion_model_cfg", {}) |
| 64 | |
| 65 | class MLLMInContext(PreTrainedModel): |
| 66 | def __init__( |