Handle dictionary in confgw

configuration_aimv2.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 
 from transformers.configuration_utils import PretrainedConfig
 
@@ -148,8 +148,8 @@ class AIMv2Config(PretrainedConfig):
 
     def __init__(
         self,
-        vision_config: Optional[AIMv2VisionConfig] = None,
-        text_config: Optional[AIMv2TextConfig] = None,
+        vision_config: Optional[Union[AIMv2VisionConfig, Dict[str, Any]]] = None,
+        text_config: Optional[Union[AIMv2TextConfig, Dict[str, Any]]] = None,
         projection_dim: int = 768,
         init_temperature: float = 0.07,
         max_logit_scale: float = 100.0,
@@ -158,8 +158,13 @@ class AIMv2Config(PretrainedConfig):
         super().__init__(**kwargs)
         if vision_config is None:
             vision_config = AIMv2VisionConfig()
+        elif isinstance(vision_config, dict):
+            vision_config = AIMv2VisionConfig(**vision_config)
+
         if text_config is None:
             text_config = AIMv2TextConfig()
+        elif isinstance(text_config, dict):
+            text_config = AIMv2TextConfig(**text_config)
 
         self.vision_config = vision_config
         self.text_config = text_config