| """HuggingFace PretrainedConfig for VisionPsyNano.""" |
| from __future__ import annotations |
|
|
| import os |
| from typing import Any, Optional |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| FLASH_MIN_SIDE_LEN = 512 |
|
|
|
|
| def resolve_is_flash( |
| *, |
| is_flash: Optional[bool] = None, |
| variant: Optional[str] = None, |
| resize_to_max_side_len: Optional[bool] = None, |
| ) -> bool: |
| """Resolve Flash preprocess mode from explicit flag, legacy variant, or resize policy.""" |
| if is_flash is not None: |
| return bool(is_flash) |
| if variant is not None: |
| v = str(variant).lower().strip() |
| if v in ("flash", "nano-flash", "visionpsy-nano-flash"): |
| return True |
| if v in ("nano", "plain", "nano-plain", "visionpsy-nano"): |
| return False |
| raise ValueError( |
| "legacy variant must be 'nano' or 'flash' " |
| f"(or aliases); got {variant!r}. Prefer is_flash=True/False." |
| ) |
| if resize_to_max_side_len is not None: |
| return not bool(resize_to_max_side_len) |
| return False |
|
|
|
|
| def apply_flash_preprocess( |
| *, |
| is_flash: bool, |
| resize_to_max_side_len: Optional[bool] = None, |
| resize_min_side_len: Optional[int] = None, |
| ) -> tuple[bool, Optional[int]]: |
| """Return (resize_to_max_side_len, resize_min_side_len) for the given Flash mode.""" |
| if resize_to_max_side_len is None: |
| resize_to_max_side_len = not is_flash |
| resize_to_max_side_len = bool(resize_to_max_side_len) |
| if is_flash: |
| resize_min_side_len = max(int(resize_min_side_len or 0), FLASH_MIN_SIDE_LEN) |
| elif resize_to_max_side_len: |
| resize_min_side_len = None |
| return resize_to_max_side_len, resize_min_side_len |
|
|
|
|
| _DEFAULT_CHAT_TEMPLATE = ( |
| "{% for message in messages %}" |
| "{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}" |
| "{% endfor %}" |
| "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}" |
| ) |
|
|
| _DEFAULT_EXTRA_TOKENS = { |
| "image_token": "<|image|>", |
| "global_image_token": "<|global_image|>", |
| "r1c1": "<row_1_col_1>", |
| "r1c2": "<row_1_col_2>", |
| "r1c3": "<row_1_col_3>", |
| "r1c4": "<row_1_col_4>", |
| "r1c5": "<row_1_col_5>", |
| "r1c6": "<row_1_col_6>", |
| "r1c7": "<row_1_col_7>", |
| "r1c8": "<row_1_col_8>", |
| "r2c1": "<row_2_col_1>", |
| "r2c2": "<row_2_col_2>", |
| "r2c3": "<row_2_col_3>", |
| "r2c4": "<row_2_col_4>", |
| "r2c5": "<row_2_col_5>", |
| "r2c6": "<row_2_col_6>", |
| "r2c7": "<row_2_col_7>", |
| "r2c8": "<row_2_col_8>", |
| "r3c1": "<row_3_col_1>", |
| "r3c2": "<row_3_col_2>", |
| "r3c3": "<row_3_col_3>", |
| "r3c4": "<row_3_col_4>", |
| "r3c5": "<row_3_col_5>", |
| "r3c6": "<row_3_col_6>", |
| "r3c7": "<row_3_col_7>", |
| "r3c8": "<row_3_col_8>", |
| "r4c1": "<row_4_col_1>", |
| "r4c2": "<row_4_col_2>", |
| "r4c3": "<row_4_col_3>", |
| "r4c4": "<row_4_col_4>", |
| "r4c5": "<row_4_col_5>", |
| "r4c6": "<row_4_col_6>", |
| "r4c7": "<row_4_col_7>", |
| "r4c8": "<row_4_col_8>", |
| "r5c1": "<row_5_col_1>", |
| "r5c2": "<row_5_col_2>", |
| "r5c3": "<row_5_col_3>", |
| "r5c4": "<row_5_col_4>", |
| "r5c5": "<row_5_col_5>", |
| "r5c6": "<row_5_col_6>", |
| "r5c7": "<row_5_col_7>", |
| "r5c8": "<row_5_col_8>", |
| "r6c1": "<row_6_col_1>", |
| "r6c2": "<row_6_col_2>", |
| "r6c3": "<row_6_col_3>", |
| "r6c4": "<row_6_col_4>", |
| "r6c5": "<row_6_col_5>", |
| "r6c6": "<row_6_col_6>", |
| "r6c7": "<row_6_col_7>", |
| "r6c8": "<row_6_col_8>", |
| "r7c1": "<row_7_col_1>", |
| "r7c2": "<row_7_col_2>", |
| "r7c3": "<row_7_col_3>", |
| "r7c4": "<row_7_col_4>", |
| "r7c5": "<row_7_col_5>", |
| "r7c6": "<row_7_col_6>", |
| "r7c7": "<row_7_col_7>", |
| "r7c8": "<row_7_col_8>", |
| "r8c1": "<row_8_col_1>", |
| "r8c2": "<row_8_col_2>", |
| "r8c3": "<row_8_col_3>", |
| "r8c4": "<row_8_col_4>", |
| "r8c5": "<row_8_col_5>", |
| "r8c6": "<row_8_col_6>", |
| "r8c7": "<row_8_col_7>", |
| "r8c8": "<row_8_col_8>", |
| } |
|
|
|
|
| class VisionPsyNanoConfig(PretrainedConfig): |
| """Config for VisionPsyNano (``is_flash`` selects preprocess mode).""" |
|
|
| model_type = "visionpsynano" |
| keys_to_ignore_at_inference = ["past_key_values"] |
|
|
| def __init__( |
| self, |
| is_flash: Optional[bool] = None, |
| variant: Optional[str] = None, |
| vit_hidden_dim: int = 768, |
| vit_inter_dim: int = 3072, |
| vit_patch_size: int = 16, |
| vit_img_size: int = 512, |
| vit_n_heads: int = 12, |
| vit_dropout: float = 0.0, |
| vit_n_blocks: int = 12, |
| vit_ln_eps: float = 1e-6, |
| vit_cls_flag: bool = False, |
| vit_model_type: str = "google/siglip2-base-patch16-512", |
| lm_hidden_dim: int = 960, |
| lm_inter_dim: int = 2560, |
| lm_rms_eps: float = 1e-5, |
| lm_re_base: int = 100000, |
| lm_max_position_embeddings: int = 8192, |
| lm_base_vocab_size: int = 49152, |
| extra_token_amount: int = 66, |
| lm_vocab_size: Optional[int] = None, |
| lm_n_heads: int = 15, |
| lm_n_kv_heads: int = 5, |
| lm_dropout: float = 0.0, |
| lm_n_blocks: int = 32, |
| lm_attn_scaling: float = 1.0, |
| lm_max_length: int = 4096, |
| lm_use_tokens: bool = False, |
| lm_tie_weights: bool = True, |
| lm_model_type: str = "HuggingFaceTB/SmolLM2-360M-Instruct", |
| lm_tokenizer: str = "HuggingFaceTB/SmolLM2-360M-Instruct", |
| lm_chat_template: str = _DEFAULT_CHAT_TEMPLATE, |
| mp_pixel_shuffle_factor: int = 4, |
| mp_image_token_length: int = 64, |
| max_img_size: int = 2048, |
| resize_to_max_side_len: Optional[bool] = None, |
| resize_min_side_len: Optional[int] = None, |
| inference_max_img_size: Optional[int] = None, |
| vlm_extra_tokens: Optional[dict] = None, |
| vlm_load_backbone_weights: bool = True, |
| vlm_checkpoint_path: str = "checkpoints", |
| hf_repo_name: str = "qvac/VisionPsy-Nano-460M", |
| compile_inference: bool = True, |
| compile_inference_mode: str = "reduce-overhead", |
| cuda_graphs_cache_quantum: int = 128, |
| eos_check_interval: int = 16, |
| **kwargs: Any, |
| ): |
| is_flash = resolve_is_flash( |
| is_flash=is_flash, |
| variant=variant, |
| resize_to_max_side_len=resize_to_max_side_len, |
| ) |
| resize_to_max_side_len, resize_min_side_len = apply_flash_preprocess( |
| is_flash=is_flash, |
| resize_to_max_side_len=resize_to_max_side_len, |
| resize_min_side_len=resize_min_side_len, |
| ) |
|
|
| if lm_vocab_size is None: |
| lm_vocab_size = lm_base_vocab_size + extra_token_amount |
|
|
| self.is_flash = bool(is_flash) |
| self.vit_hidden_dim = vit_hidden_dim |
| self.vit_inter_dim = vit_inter_dim |
| self.vit_patch_size = vit_patch_size |
| self.vit_img_size = vit_img_size |
| self.vit_n_heads = vit_n_heads |
| self.vit_dropout = vit_dropout |
| self.vit_n_blocks = vit_n_blocks |
| self.vit_ln_eps = vit_ln_eps |
| self.vit_cls_flag = vit_cls_flag |
| self.vit_model_type = vit_model_type |
|
|
| self.lm_hidden_dim = lm_hidden_dim |
| self.lm_inter_dim = lm_inter_dim |
| self.lm_rms_eps = lm_rms_eps |
| self.lm_re_base = lm_re_base |
| self.lm_max_position_embeddings = lm_max_position_embeddings |
| self.lm_base_vocab_size = lm_base_vocab_size |
| self.extra_token_amount = extra_token_amount |
| self.lm_vocab_size = lm_vocab_size |
| self.lm_n_heads = lm_n_heads |
| self.lm_n_kv_heads = lm_n_kv_heads |
| self.lm_dropout = lm_dropout |
| self.lm_n_blocks = lm_n_blocks |
| self.lm_attn_scaling = lm_attn_scaling |
| self.lm_max_length = lm_max_length |
| self.lm_use_tokens = lm_use_tokens |
| self.lm_tie_weights = lm_tie_weights |
| self.lm_model_type = lm_model_type |
| self.lm_tokenizer = lm_tokenizer |
| self.lm_chat_template = lm_chat_template |
|
|
| self.mp_pixel_shuffle_factor = mp_pixel_shuffle_factor |
| self.mp_image_token_length = mp_image_token_length |
| self.max_img_size = max_img_size |
| self.resize_to_max_side_len = bool(resize_to_max_side_len) |
| self.resize_min_side_len = resize_min_side_len |
| self.inference_max_img_size = inference_max_img_size |
| self.vlm_extra_tokens = dict(vlm_extra_tokens or _DEFAULT_EXTRA_TOKENS) |
| self.vlm_load_backbone_weights = vlm_load_backbone_weights |
| self.vlm_checkpoint_path = vlm_checkpoint_path |
| if self.is_flash and hf_repo_name == "qvac/VisionPsy-Nano-460M": |
| self.hf_repo_name = "qvac/VisionPsy-Nano-460M-Flash" |
| else: |
| self.hf_repo_name = hf_repo_name |
|
|
| self.compile_inference = compile_inference |
| self.compile_inference_mode = compile_inference_mode |
| self.cuda_graphs_cache_quantum = cuda_graphs_cache_quantum |
| self.eos_check_interval = eos_check_interval |
|
|
| kwargs.pop("text_config", None) |
| kwargs.pop("vision_config", None) |
|
|
| super().__init__(**kwargs) |
|
|
| def get_text_config(self, decoder: bool = False, **kwargs): |
| """VisionPsyNano is a single flat config (not text+vision composite).""" |
| return self |
|
|
| @property |
| def variant(self) -> str: |
| return "flash" if self.is_flash else "nano" |
|
|
| def to_vlm_config(self): |
| """Convert to the internal VLMConfig dataclass used by the core modules.""" |
| from dataclasses import fields as dc_fields |
|
|
| try: |
| from .vlm_config import VLMConfig |
| except ImportError: |
| from vlm_config import VLMConfig |
|
|
| valid = {f.name for f in dc_fields(VLMConfig)} |
| payload = {k: getattr(self, k) for k in valid if hasattr(self, k)} |
| return VLMConfig(**payload) |
|
|
| @classmethod |
| def from_vlm_config( |
| cls, cfg, *, is_flash: Optional[bool] = None, variant: Optional[str] = None |
| ) -> "VisionPsyNanoConfig": |
| from dataclasses import asdict |
|
|
| data = asdict(cfg) |
| data["is_flash"] = resolve_is_flash( |
| is_flash=is_flash, |
| variant=variant, |
| resize_to_max_side_len=data.get("resize_to_max_side_len"), |
| ) |
| data.pop("variant", None) |
| return cls(**data) |
|
|
| @classmethod |
| def from_legacy_dict( |
| cls, |
| raw: dict, |
| *, |
| is_flash: Optional[bool] = None, |
| variant: Optional[str] = None, |
| ) -> "VisionPsyNanoConfig": |
| """Load an existing VisionPsyNano / nanoVLM config.json (without model_type).""" |
| raw = dict(raw) |
| for key in ( |
| "model_type", |
| "architectures", |
| "auto_map", |
| "transformers_version", |
| "text_config", |
| "vision_config", |
| ): |
| raw.pop(key, None) |
| raw["is_flash"] = resolve_is_flash( |
| is_flash=is_flash if is_flash is not None else raw.get("is_flash"), |
| variant=variant if variant is not None else raw.get("variant"), |
| resize_to_max_side_len=raw.get("resize_to_max_side_len"), |
| ) |
| raw.pop("variant", None) |
| return cls(**raw) |
|
|
