
    Z j@                     p    S SK Jr  SSKJr  SSKJr  SSKJrJr  \" SS9\ " S	 S
\5      5       5       r	S
/r
g)    )strict   )PreTrainedConfig)auto_docstring   )CONFIG_MAPPING
AutoConfigzlerobot/pi0_base)
checkpointc                     ^  \ rS rSr% SrSr\\S.rSr\	\
-  S-  \S'   Sr\	\
-  S-  \S'   Sr\\S	'   S
r\\S'   S
r\\S'   Sr\\S'   Sr\\S'   Sr\\S'   Sr\\S'   Sr\\S'   Sr\\S'   Sr\\S'   Sr\\S'   U 4S jrS rSrU =r$ ) 	PI0Config   aX  
vlm_config (`dict`, *optional*):
    Configuration for the vlm backbone (PaliGemmaModel).
dit_config (`dict`, *optional*):
    Configuration for the DiT backbone. Defaults to a Gemma 300M variant.
chunk_size (`int`, *optional*, defaults to 50):
    Number of action steps to predict per chunk.
max_state_dim (`int`, *optional*, defaults to 32):
    Maximum state vector dimension (shorter vectors are zero-padded).
max_action_dim (`int`, *optional*, defaults to 32):
    Maximum action vector dimension (shorter vectors are zero-padded).
num_inference_steps (`int`, *optional*, defaults to 10):
    Number of denoising steps during inference.
time_sampling_beta_alpha (`float`, *optional*, defaults to 1.5):
    Alpha parameter for Beta distribution used to sample diffusion time during training.
time_sampling_beta_beta (`float`, *optional*, defaults to 1.0):
    Beta parameter for Beta distribution used to sample diffusion time during training.
time_sampling_scale (`float`, *optional*, defaults to 0.999):
    Scale factor for sampled time values.
time_sampling_offset (`float`, *optional*, defaults to 0.001):
    Offset added to sampled time values.
min_period (`float`, *optional*, defaults to 0.004):
    Minimum period for sinusoidal time embedding.
max_period (`float`, *optional*, defaults to 4.0):
    Maximum period for sinusoidal time embedding.
loss_reduction (`str`, *optional*, defaults to `"mean"`):
    The reduction to use on MSE loss.

Example:
```python
>>> from transformers import PI0ForConditionalGeneration, PI0Config

>>> config = PI0Config()
>>> model = PI0ForConditionalGeneration(config)
```
pi0)
vlm_config
dit_configNr   r   2   
chunk_size    max_state_dimmax_action_dim
   num_inference_stepsg      ?time_sampling_beta_alphag      ?time_sampling_beta_betag+?time_sampling_scalegMbP?time_sampling_offsetgMbp?
min_periodg      @
max_periodmeanloss_reductionc                   > [        U R                  [        5      (       a;  U R                  R                  SS5      n[        U   " S0 U R                  D6U l        O4U R                  c'  [        S   " SSSSSSS	S
.SSSSSSSS	SS.	SS	S9U l        [        U R
                  [        5      (       a;  U R
                  R                  SS5      n[        U   " S0 U R
                  D6U l        OCU R
                  c6  [        S   " SSSSSSU R                  R                  R                  S9U l        SU R
                  l        SU R
                  l	        SU R                  R                  l	        [        TU ],  " S0 UD6  g )N
model_type	paligemmagemmai      i @        i )r!   hidden_sizenum_hidden_layersintermediate_sizenum_attention_headsnum_key_value_heads
vocab_sizesiglip_vision_modeli  i        r      F)	r!   r)   r'   
patch_size
image_sizer(   r*   r,   vision_use_head)text_configvision_configprojection_dimimage_token_idi   i      )r'   r(   r)   r*   r+   head_dimr,   T )
isinstancer   dictgetr   r   r4   r,   	is_causaluse_bidirectional_attentionsuper__post_init__)selfkwargsvlm_model_typedit_model_type	__class__s       z/root/GenerationalWealth/GenerationalWealth/venv/lib/python3.13/site-packages/transformers/models/pi0/configuration_pi0.pyrA   PI0Config.__post_init__T   s^   doot,,!__00{KN,^<OtODO__$,[9")#')+).+,+,"( #8)-#'"$"%)++-"(',
  $%-DO2 doot,,!__00wGN,^<OtODO__$,W5 "$"&$%$%??66AADO %*!6:3BF##?''    c                     U R                   R                  S-  S:w  a-  [        SU R                  R                   R                   S35      eg)zOPart of `@strict`-powered validation. Validates the architecture of the config.r   r   zDiT hidden dim=(z) must be divisible by 2N)r   r'   
ValueErrorconfig)rB   s    rG   validate_architecturePI0Config.validate_architecture   sE    ??&&*a//0F0F0R0R/SSklmm 0rI   )r   r   )__name__
__module____qualname____firstlineno____doc__r!   r	   sub_configsr   r<   r   __annotations__r   r   intr   r   r   r   floatr   r   r   r   r   r   strrA   rM   __static_attributes____classcell__)rF   s   @rG   r   r      s    #J J!+:FK15J''$.515J''$.5JM3NC!!&)e)%(U(!&&"'%'JJ NC 0(dn nrI   r   N)huggingface_hub.dataclassesr   configuration_utilsr   utilsr   autor   r	   r   __all__r:   rI   rG   <module>r`      sM   ( / 3 # - -.ln  ln  /ln^ -rI   