update readme and minor fix

Author: bfshi

README.md

@@ -32,7 +32,7 @@
 <hr style="border: 2px solid gray;"></hr>
 
 ## Latest Updates
-- [2025.5.21] Models & code of PS3 and VILA-HD are released! We released two PS3 models (`PS3-1.5K-SigLIP` and `PS3-4K-SigLIP`) and two VILA-HD models (`VILA-HD-1.5K-8B-SigLIP` and `VILA-HD-4K-8B-SigLIP`).
+- [2025.6.4] Models & code of PS3 and VILA-HD are released! We released two PS3 models (`PS3-1.5K-SigLIP` and `PS3-4K-SigLIP`) and two VILA-HD models (`VILA-HD-1.5K-8B-SigLIP` and `VILA-HD-4K-8B-SigLIP`), and the corresponding training/inference code are also released.
 - [2025.4.22] Demo of VILA-HD is released! Welcome to give it a try. We are actively improving the model so any feedback is welcome!
 - [2025.4.4] Selected as conference highlight at CVPR 2025. See you in Nashville!
 - [2025.3.24] Initial paper release. Code and weights of PS3 and VILA-HD will be released very soon!
@@ -317,7 +317,7 @@ class PS3VisionModel(PS3PreTrainedModel):
 
 ## Training
 
-Coming soon.
+Please see `train/`.
 
 ## Using PS3 in Downstream MLLMs
 

train/src/open_clip/factory.py

@@ -402,7 +402,7 @@ def _convert_ln(m):
 
         if checkpoint_path:
             logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
-            load_checkpoint(model, checkpoint_path, strict=False, weights_only=load_weights_only)
+            load_checkpoint(model, checkpoint_path, strict=True, weights_only=load_weights_only)
         else:
             error_str = (
                 f'Pretrained weights ({pretrained}) not found for model {model_name}.'

train/src/open_clip/model.py

@@ -470,10 +470,7 @@ def __init__(
 
         self.ps3 = vision_cfg.get("ps3", False) if isinstance(vision_cfg, dict) else vision_cfg.ps3
         if self.ps3:
-            if "ps3_feature_proj" in vision_cfg and vision_cfg["ps3_feature_proj"] == "mlp":
-                self.prompt_proj = Mlp(in_features=self.text.output_dim, hidden_features=self.visual.width, out_features=self.visual.width, norm_layer=nn.LayerNorm)
-            else:
-                self.prompt_proj = nn.Linear(self.text.output_dim, self.visual.width)
+            self.prompt_proj = Mlp(in_features=self.text.output_dim, hidden_features=self.visual.width, out_features=self.visual.width, norm_layer=nn.LayerNorm)
 
     def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
         # lock image tower as per LiT - https://arxiv.org/abs/2111.07991