frappe/models/utils.py at main · OpenHelix-Team/frappe · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import os
from torchvision.datasets.utils import download_url
import torch
import torchvision.models as torchvision_models
from typing import Any, Optional
import math
import warnings

@torch.no_grad()
def load_encoders(enc_type, device, resolution=256):
    enc_names = enc_type.split(',')
    encoders, architectures, encoder_types = [], [], []
    for enc_name in enc_names:
        encoder_type, architecture, model_config = enc_name.split('-')
        architectures.append(architecture)
        encoder_types.append(encoder_type)
        if 'dinov2' in encoder_type:
            import timm
            dinov2_path = os.path.expanduser('path_to/facebookresearch_dinov2_main')

            model_name = f'dinov2_vit{model_config}14'
            encoder = torch.hub.load(dinov2_path, model_name, source='local', pretrained=False)

            checkpoint_name = f'{model_name}_pretrain.pth'
            checkpoint_path = os.path.expanduser(f'/.cache/torch/hub/checkpoints/{checkpoint_name}')

            if os.path.exists(checkpoint_path):
                state_dict = torch.load(checkpoint_path, map_location="cpu")
                encoder.load_state_dict(state_dict, strict=True)
            else:
                print(f"Warning: DINOv2 checkpoint not found at {checkpoint_path}. Model is NOT pretrained.")

            del encoder.head
            patch_resolution = 16 * (resolution // 256)
            encoder.pos_embed.data = timm.layers.pos_embed.resample_abs_pos_embed(
                encoder.pos_embed.data, [patch_resolution, patch_resolution],
            )
            encoder.head = torch.nn.Identity()
            encoder = encoder.to(device)
            encoder.eval()


        elif encoder_type == 'clip':
            from transformers import CLIPVisionModel
            model_id = "path_to/CLIP-ViT-H-14-laion2B-s32B-b79K"
            encoder = CLIPVisionModel.from_pretrained(model_id)
            encoder.embed_dim = encoder.config.hidden_size #1280
            encoder.forward_features = lambda pixel_values: encoder.forward(
                    pixel_values=pixel_values
                        ).last_hidden_state[:, 1:, :]
            encoder = encoder.to(device)
            encoder.eval()

        elif encoder_type == 'theia':
            from transformers import AutoModel

            local_theia_path = "path_to/theia-base-patch16-224-cdiv"
            encoder = AutoModel.from_pretrained(
                    local_theia_path,
                    trust_remote_code=True
                )

            encoder.embed_dim = encoder.backbone.model.config.hidden_size

            def patched_theia_forward(x, feature_reduction_method='none'):

                input_for_backbone = x.to(torch.float32)
                feature = encoder.backbone(input_for_backbone)
                if feature_reduction_method == 'none':
                    processed_output = handle_feature_output(feature, feature_reduce_method=None, num_discard_tokens=encoder.num_reg_tokens)
                else:
                    processed_output = handle_feature_output(feature, feature_reduce_method=feature_reduction_method, num_discard_tokens=encoder.num_reg_tokens)

                return processed_output.to(x.dtype)

            encoder.forward_features = patched_theia_forward
            encoder = encoder.to(device)
            encoder.eval()

        elif encoder_type == 'vit':
            from transformers import ViTModel
            model_id = "path_to/vit-huge-patch14-224-in21k"
            print(f"Loading ViT-Huge model (original output) from local path: {model_id}")
            encoder = ViTModel.from_pretrained(model_id)
            encoder.embed_dim = encoder.config.hidden_size #1280
            encoder.forward_features = lambda pixel_values: encoder.forward(pixel_values=pixel_values).last_hidden_state[:, 1:, :]
            encoder = encoder.to(device)
            encoder.eval()

        encoders.append(encoder)
    return encoders, encoder_types, architectures


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


def load_legacy_checkpoints(state_dict, encoder_depth):
    new_state_dict = dict()
    for key, value in state_dict.items():
        if 'decoder_blocks' in key:
            parts =key.split('.')
            new_idx = int(parts[1]) + encoder_depth
            parts[0] = 'blocks'
            parts[1] = str(new_idx)
            new_key = '.'.join(parts)
            new_state_dict[new_key] = value
        else:
            new_state_dict[key] = value
    return new_state_dict

#theia reuse to correctly handle feature output
def handle_feature_output(
    x: torch.Tensor, feature_reduce_method: Optional[str] = None, num_discard_tokens: int = 0
) -> torch.Tensor:
    match feature_reduce_method:
        case "mean_pooling":
            return torch.mean(x[:, 1 : x.size(1) - num_discard_tokens], dim=1)
        case "max_pooling":
            return torch.amax(x[:, 1 : x.size(1) - num_discard_tokens], dim=1)
        case "cls":
            return x[:, 0]
        case "identity":
            return x
        case None:
            return x[:, 1 : x.size(1) - num_discard_tokens]
        case _:
            raise NotImplementedError(f"feature_reduce_method {feature_reduce_method} it not implemented.")