Class CLIPModel

wan/models/wan_image_encoder.py:503–553

class CLIPModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):

    def __init__(self):
        super(CLIPModel, self).__init__()
        # init model
        self.model, self.transforms = clip_xlm_roberta_vit_h_14(
            pretrained=False,
            return_transforms=True,
            return_tokenizer=False)

    def forward(self, videos):
        # preprocess
        size = (self.model.image_size,) * 2
        videos = torch.cat([
            F.interpolate(
                u.transpose(0, 1),
                size=size,
                mode='bicubic',
                align_corners=False) for u in videos
        ])
        videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))

        # forward
        with torch.cuda.amp.autocast(dtype=self.dtype):
            out = self.model.visual(videos, use_31_block=True)
            return out

    @classmethod
    def from_pretrained(cls, pretrained_model_path, transformer_additional_kwargs={}):
        def filter_kwargs(cls, kwargs):
            import inspect
            sig = inspect.signature(cls.__init__)
            valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
            return filtered_kwargs

        model = cls(**filter_kwargs(cls, transformer_additional_kwargs))
        if pretrained_model_path.endswith(".safetensors"):
            from safetensors.torch import load_file, safe_open
            state_dict = load_file(pretrained_model_path)
        else:
            state_dict = torch.load(pretrained_model_path, map_location="cpu")
        tmp_state_dict = {} 
        for key in state_dict:
            tmp_state_dict["model." + key] = state_dict[key]
        state_dict = tmp_state_dict
        m, u = model.load_state_dict(state_dict)
        
        print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
        print(m, u)
        return model