1	import sys
2	import os.path
3	import subprocess
4	
5	custom_nodes_path = os.path.dirname(os.path.abspath(__file__))
6	
7	def build_pip_install_cmds(args):
8	    if "python_embeded" in sys.executable or "python_embedded" in sys.executable:

11	        return [sys.executable, '-m', 'pip', 'install'] + args
12	
13	def ensure_package():
14	    cmds = build_pip_install_cmds(['-r', 'requirements.txt'])
15	    subprocess.run(cmds, cwd=custom_nodes_path)
16	
17	if __name__ == "__main__":
18	    ensure_package()

187	    def __call__(self, img: PIL.Image.Image, target: dict):
188	        init_boxes = len(target["boxes"])
189	        max_patience = 10
190	        for i in range(max_patience):
191	            w = random.randint(self.min_size, min(img.width, self.max_size))
192	            h = random.randint(self.min_size, min(img.height, self.max_size))
193	            region = T.RandomCrop.get_params(img, [h, w])
194	            result_img, result_target = crop(img, target, region)

188	        init_boxes = len(target["boxes"])
189	        max_patience = 10
190	        for i in range(max_patience):
191	            w = random.randint(self.min_size, min(img.width, self.max_size))
192	            h = random.randint(self.min_size, min(img.height, self.max_size))
193	            region = T.RandomCrop.get_params(img, [h, w])
194	            result_img, result_target = crop(img, target, region)
195	            if (

217	    def __init__(self, p=0.5):
218	        self.p = p
219	
220	    def __call__(self, img, target):
221	        if random.random() < self.p:
222	            return hflip(img, target)
223	        return img, target
224	

224	
225	
226	class RandomResize(object):
227	    def __init__(self, sizes, max_size=None):
228	        assert isinstance(sizes, (list, tuple))
229	        self.sizes = sizes
230	        self.max_size = max_size
231	

229	        self.sizes = sizes
230	        self.max_size = max_size
231	
232	    def __call__(self, img, target=None):
233	        size = random.choice(self.sizes)
234	        return resize(img, target, size, self.max_size)
235	
236	

238	    def __init__(self, max_pad):
239	        self.max_pad = max_pad
240	
241	    def __call__(self, img, target):
242	        pad_x = random.randint(0, self.max_pad)
243	        pad_y = random.randint(0, self.max_pad)
244	        return pad(img, target, (pad_x, pad_y))
245	

239	        self.max_pad = max_pad
240	
241	    def __call__(self, img, target):
242	        pad_x = random.randint(0, self.max_pad)
243	        pad_y = random.randint(0, self.max_pad)
244	        return pad(img, target, (pad_x, pad_y))
245	
246	

255	        self.transforms2 = transforms2
256	        self.p = p
257	
258	    def __call__(self, img, target):
259	        if random.random() < self.p:
260	            return self.transforms1(img, target)
261	        return self.transforms2(img, target)
262	

108	        xs = self.body(tensor_list.tensors)
109	        out: Dict[str, NestedTensor] = {}
110	        for name, x in xs.items():
111	            m = tensor_list.mask
112	            assert m is not None
113	            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
114	            out[name] = NestedTensor(x, mask)
115	        # import ipdb; ipdb.set_trace()

135	            )
136	        else:
137	            raise NotImplementedError("Why you can get here with name {}".format(name))
138	        # num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
139	        assert name not in ("resnet18", "resnet34"), "Only resnet50 and resnet101 are available."
140	        assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]]
141	        num_channels_all = [256, 512, 1024, 2048]
142	        num_channels = num_channels_all[4 - len(return_interm_indices) :]

136	        else:
137	            raise NotImplementedError("Why you can get here with name {}".format(name))
138	        # num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
139	        assert name not in ("resnet18", "resnet34"), "Only resnet50 and resnet101 are available."
140	        assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]]
141	        num_channels_all = [256, 512, 1024, 2048]
142	        num_channels = num_channels_all[4 - len(return_interm_indices) :]
143	        super().__init__(backbone, train_backbone, num_channels, return_interm_indices)

174	    train_backbone = True
175	    if not train_backbone:
176	        raise ValueError("Please set lr_backbone > 0")
177	    return_interm_indices = args.return_interm_indices
178	    assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]]
179	    args.backbone_freeze_keywords
180	    use_checkpoint = getattr(args, "use_checkpoint", False)
181	

207	        bb_num_channels = backbone.num_features[4 - len(return_interm_indices) :]
208	    else:
209	        raise NotImplementedError("Unknown backbone {}".format(args.backbone))
210	
211	    assert len(bb_num_channels) == len(
212	        return_interm_indices
213	    ), f"len(bb_num_channels) {len(bb_num_channels)} != len(return_interm_indices) {len(return_interm_indices)}"
214	
215	    model = Joiner(backbone, position_embedding)
216	    model.num_channels = bb_num_channels

213	    ), f"len(bb_num_channels) {len(bb_num_channels)} != len(return_interm_indices) {len(return_interm_indices)}"
214	
215	    model = Joiner(backbone, position_embedding)
216	    model.num_channels = bb_num_channels
217	    assert isinstance(
218	        bb_num_channels, List
219	    ), "bb_num_channels is expected to be a List but {}".format(type(bb_num_channels))
220	    # import ipdb; ipdb.set_trace()
221	    return model

46	
47	    def forward(self, tensor_list: NestedTensor):
48	        x = tensor_list.tensors
49	        mask = tensor_list.mask
50	        assert mask is not None
51	        not_mask = ~mask
52	        y_embed = not_mask.cumsum(1, dtype=torch.float32)
53	        x_embed = not_mask.cumsum(2, dtype=torch.float32)

97	
98	    def forward(self, tensor_list: NestedTensor):
99	        x = tensor_list.tensors
100	        mask = tensor_list.mask
101	        assert mask is not None
102	        not_mask = ~mask
103	        y_embed = not_mask.cumsum(1, dtype=torch.float32)
104	        x_embed = not_mask.cumsum(2, dtype=torch.float32)

211	        self.num_heads = num_heads
212	        self.window_size = window_size
213	        self.shift_size = shift_size
214	        self.mlp_ratio = mlp_ratio
215	        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
216	
217	        self.norm1 = norm_layer(dim)
218	        self.attn = WindowAttention(

243	            mask_matrix: Attention mask for cyclic shift.
244	        """
245	        B, L, C = x.shape
246	        H, W = self.H, self.W
247	        assert L == H * W, "input feature has wrong size"
248	
249	        shortcut = x
250	        x = self.norm1(x)

317	            x: Input feature, tensor size (B, H*W, C).
318	            H, W: Spatial resolution of the input feature.
319	        """
320	        B, L, C = x.shape
321	        assert L == H * W, "input feature has wrong size"
322	
323	        x = x.view(B, H, W, C)
324	

746	        # collect for nesttensors
747	        outs_dict = {}
748	        for idx, out_i in enumerate(outs):
749	            m = tensor_list.mask
750	            assert m is not None
751	            mask = F.interpolate(m[None].float(), size=out_i.shape[-2:]).to(torch.bool)[0]
752	            outs_dict[idx] = NestedTensor(out_i, mask)
753	

759	        self._freeze_stages()
760	
761	
762	def build_swin_transformer(modelname, pretrain_img_size, **kw):
763	    assert modelname in [
764	        "swin_T_224_1k",
765	        "swin_B_224_22k",
766	        "swin_B_384_22k",
767	        "swin_L_224_22k",
768	        "swin_L_384_22k",
769	    ]
770	
771	    model_para_dict = {
772	        "swin_T_224_1k": dict(

105	        self.head_dim = embed_dim // num_heads
106	        self.v_dim = v_dim
107	        self.l_dim = l_dim
108	
109	        assert (
110	            self.head_dim * self.num_heads == self.embed_dim
111	        ), f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
112	        self.scale = self.head_dim ** (-0.5)
113	        self.dropout = dropout
114	

83	        self.sub_sentence_present = sub_sentence_present
84	
85	        # setting query dim
86	        self.query_dim = query_dim
87	        assert query_dim == 4
88	
89	        # for dn training
90	        self.num_patterns = num_patterns

129	                )
130	                in_channels = hidden_dim
131	            self.input_proj = nn.ModuleList(input_proj_list)
132	        else:
133	            assert two_stage_type == "no", "two_stage_type should be no if num_feature_levels=1 !!!"
134	            self.input_proj = nn.ModuleList(
135	                [
136	                    nn.Sequential(

144	        self.aux_loss = aux_loss
145	        self.box_pred_damping = box_pred_damping = None
146	
147	        self.iter_update = iter_update
148	        assert iter_update, "Why not iter_update?"
149	
150	        # prepare pred layers
151	        self.dec_pred_bbox_embed_share = dec_pred_bbox_embed_share

169	        self.transformer.decoder.class_embed = self.class_embed
170	
171	        # two stage
172	        self.two_stage_type = two_stage_type
173	        assert two_stage_type in ["no", "standard"], "unknown param {} of two_stage_type".format(
174	            two_stage_type
175	        )
176	        if two_stage_type != "no":
177	            if two_stage_bbox_embed_share:
178	                assert dec_pred_bbox_embed_share

174	            two_stage_type
175	        )
176	        if two_stage_type != "no":
177	            if two_stage_bbox_embed_share:
178	                assert dec_pred_bbox_embed_share
179	                self.transformer.enc_out_bbox_embed = _bbox_embed
180	            else:
181	                self.transformer.enc_out_bbox_embed = copy.deepcopy(_bbox_embed)

180	            else:
181	                self.transformer.enc_out_bbox_embed = copy.deepcopy(_bbox_embed)
182	
183	            if two_stage_class_embed_share:
184	                assert dec_pred_bbox_embed_share
185	                self.transformer.enc_out_class_embed = _class_embed
186	            else:
187	                self.transformer.enc_out_class_embed = copy.deepcopy(_class_embed)

283	        for l, feat in enumerate(features):
284	            src, mask = feat.decompose()
285	            srcs.append(self.input_proj[l](src))
286	            masks.append(mask)
287	            assert mask is not None
288	        if self.num_feature_levels > len(srcs):
289	            _len_srcs = len(srcs)
290	            for l in range(_len_srcs, self.num_feature_levels):

224	
225	        bs, num_query, _ = query.shape
226	        bs, num_value, _ = value.shape
227	
228	        assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value
229	
230	        value = self.value_proj(value)
231	        if key_padding_mask is not None:

77	        self.num_encoder_layers = num_encoder_layers
78	        self.num_unicoder_layers = num_unicoder_layers
79	        self.num_decoder_layers = num_decoder_layers
80	        self.num_queries = num_queries
81	        assert query_dim == 4
82	
83	        # choose encoder layer type
84	        encoder_layer = DeformableTransformerEncoderLayer(

107	        else:
108	            feature_fusion_layer = None
109	
110	        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
111	        assert encoder_norm is None
112	        self.encoder = TransformerEncoder(
113	            encoder_layer,
114	            num_encoder_layers,

158	            else:
159	                self.level_embed = None
160	
161	        self.learnable_tgt_init = learnable_tgt_init
162	        assert learnable_tgt_init, "why not learnable_tgt_init"
163	        self.embed_init_tgt = embed_init_tgt
164	        if (two_stage_type != "no" and embed_init_tgt) or (two_stage_type == "no"):
165	            self.tgt_embed = nn.Embedding(self.num_queries, d_model)

168	            self.tgt_embed = None
169	
170	        # for two stage
171	        self.two_stage_type = two_stage_type
172	        assert two_stage_type in ["no", "standard"], "unknown param {} of two_stage_type".format(
173	            two_stage_type
174	        )
175	        if two_stage_type == "standard":
176	            # anchor selection at the output of encoder
177	            self.enc_output = nn.Linear(d_model, d_model)

613	            self.layers = []
614	        self.num_layers = num_layers
615	        self.norm = norm
616	        self.return_intermediate = return_intermediate
617	        assert return_intermediate, "support return_intermediate only"
618	        self.query_dim = query_dim
619	        assert query_dim in [2, 4], "query_dim should be 2/4 but {}".format(query_dim)
620	        self.num_feature_levels = num_feature_levels

615	        self.norm = norm
616	        self.return_intermediate = return_intermediate
617	        assert return_intermediate, "support return_intermediate only"
618	        self.query_dim = query_dim
619	        assert query_dim in [2, 4], "query_dim should be 2/4 but {}".format(query_dim)
620	        self.num_feature_levels = num_feature_levels
621	
622	        self.ref_point_head = MLP(query_dim // 2 * d_model, d_model, d_model, 2)

669	                    reference_points[:, :, None]
670	                    * torch.cat([valid_ratios, valid_ratios], -1)[None, :]
671	                )  # nq, bs, nlevel, 4
672	            else:
673	                assert reference_points.shape[-1] == 2
674	                reference_points_input = reference_points[:, :, None] * valid_ratios[None, :]
675	            query_sine_embed = gen_sineembed_for_position(
676	                reference_points_input[:, :, 0, :]

845	        self.norm3 = nn.LayerNorm(d_model)
846	
847	        self.key_aware_proj = None
848	        self.use_text_feat_guide = use_text_feat_guide
849	        assert not use_text_feat_guide
850	        self.use_text_cross_attention = use_text_cross_attention
851	
852	    def rm_self_attn_modules(self):

889	        Input:
890	            - tgt/tgt_query_pos: nq, bs, d_model
891	            -
892	        """
893	        assert cross_attn_mask is None
894	
895	        # self attention
896	        if self.self_attn is not None:

252	            }
253	        Returns:
254	            _type_: _description_
255	        """
256	        assert isinstance(text_dict, dict)
257	
258	        y = text_dict["encoded_text"]
259	        text_token_mask = text_dict["text_token_mask"]

11	def build_model(args):
12	    # we use register to maintain models from catdet6 on.
13	    from .registry import MODULE_BUILD_FUNCS
14	
15	    assert args.modelname in MODULE_BUILD_FUNCS._module_dict
16	    build_func = MODULE_BUILD_FUNCS.get(args.modelname)
17	    model = build_func(args)
18	    return model

46	    and M = len(boxes2)
47	    """
48	    # degenerate boxes gives inf / nan results
49	    # so do an early check
50	    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
51	    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
52	    # except:
53	    #     import ipdb; ipdb.set_trace()

47	    """
48	    # degenerate boxes gives inf / nan results
49	    # so do an early check
50	    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
51	    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
52	    # except:
53	    #     import ipdb; ipdb.set_trace()
54	    iou, union = box_iou(boxes1, boxes2)

89	        - giou: N, 4
90	    """
91	    # degenerate boxes gives inf / nan results
92	    # so do an early check
93	    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
94	    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
95	    assert boxes1.shape == boxes2.shape
96	    iou, union = box_iou_pairwise(boxes1, boxes2)  # N, 4

90	    """
91	    # degenerate boxes gives inf / nan results
92	    # so do an early check
93	    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
94	    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
95	    assert boxes1.shape == boxes2.shape
96	    iou, union = box_iou_pairwise(boxes1, boxes2)  # N, 4
97	

91	    # degenerate boxes gives inf / nan results
92	    # so do an early check
93	    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
94	    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
95	    assert boxes1.shape == boxes2.shape
96	    iou, union = box_iou_pairwise(boxes1, boxes2)  # N, 4
97	
98	    lt = torch.min(boxes1[:, :2], boxes2[:, :2])

9	import functools
10	import io
11	import json
12	import os
13	import pickle
14	import subprocess
15	import time
16	from collections import OrderedDict, defaultdict, deque

10	import io
11	import json
12	import os
13	import pickle
14	import subprocess
15	import time
16	from collections import OrderedDict, defaultdict, deque
17	from typing import List, Optional

142	        print("gathering on cpu")
143	        dist.all_gather(size_list, local_size, group=cpu_group)
144	    size_list = [int(size.item()) for size in size_list]
145	    max_size = max(size_list)
146	    assert isinstance(local_size.item(), int)
147	    local_size = int(local_size.item())
148	
149	    # receiving Tensor from all ranks

252	    def update(self, **kwargs):
253	        for k, v in kwargs.items():
254	            if isinstance(v, torch.Tensor):
255	                v = v.item()
256	            assert isinstance(v, (float, int))
257	            self.meters[k].update(v)
258	
259	    def __getattr__(self, attr):

362	def get_sha():
363	    cwd = os.path.dirname(os.path.abspath(__file__))
364	
365	    def _run(command):
366	        return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
367	
368	    sha = "N/A"
369	    diff = "clean"

369	    diff = "clean"
370	    branch = "N/A"
371	    try:
372	        sha = _run(["git", "rev-parse", "HEAD"])
373	        subprocess.check_output(["git", "diff"], cwd=cwd)
374	        diff = _run(["git", "diff-index", "HEAD"])
375	        diff = "has uncommited changes" if diff else "clean"
376	        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])

369	    diff = "clean"
370	    branch = "N/A"
371	    try:
372	        sha = _run(["git", "rev-parse", "HEAD"])
373	        subprocess.check_output(["git", "diff"], cwd=cwd)
374	        diff = _run(["git", "diff-index", "HEAD"])
375	        diff = "has uncommited changes" if diff else "clean"
376	        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])

373	        subprocess.check_output(["git", "diff"], cwd=cwd)
374	        diff = _run(["git", "diff-index", "HEAD"])
375	        diff = "has uncommited changes" if diff else "clean"
376	        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
377	    except Exception:
378	        pass
379	    message = f"sha: {sha}, status: {diff}, branch: {branch}"
380	    return message
381	

426	        # type: (Device) -> NestedTensor # noqa
427	        cast_tensor = self.tensors.to(device)
428	        mask = self.mask
429	        if mask is not None:
430	            assert mask is not None
431	            cast_mask = mask.to(device)
432	        else:
433	            cast_mask = None

433	            cast_mask = None
434	        return NestedTensor(cast_tensor, cast_mask)
435	
436	    def to_img_list_single(self, tensor, mask):
437	        assert tensor.dim() == 3, "dim of tensor should be 3 but {}".format(tensor.dim())
438	        maxH = (~mask).sum(0).max()
439	        maxW = (~mask).sum(1).max()
440	        img = tensor[:, :maxH, :maxW]

2	# Modified from mmcv
3	# ==========================================================
4	
5	import json
6	import pickle
7	from abc import ABCMeta, abstractmethod
8	from pathlib import Path
9	

39	    img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
40	) -> torch.FloatTensor:
41	    # img: tensor(3,H,W) or tensor(B,3,H,W)
42	    # return: same as img
43	    assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim()
44	    if img.dim() == 3:
45	        assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % (
46	            img.size(0),

41	    # img: tensor(3,H,W) or tensor(B,3,H,W)
42	    # return: same as img
43	    assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim()
44	    if img.dim() == 3:
45	        assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % (
46	            img.size(0),
47	            str(img.size()),
48	        )
49	        img_perm = img.permute(1, 2, 0)
50	        mean = torch.Tensor(mean)
51	        std = torch.Tensor(std)

51	        std = torch.Tensor(std)
52	        img_res = img_perm * std + mean
53	        return img_res.permute(2, 0, 1)
54	    else:  # img.dim() == 4
55	        assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % (
56	            img.size(1),
57	            str(img.size()),
58	        )
59	        img_perm = img.permute(0, 2, 3, 1)
60	        mean = torch.Tensor(mean)
61	        std = torch.Tensor(std)

283	        """
284	        Input:
285	            pred, gt: Tensor()
286	        """
287	        assert pred.shape == gt.shape
288	        self.tp += torch.logical_and(pred == 1, gt == 1).sum().item()
289	        self.fp += torch.logical_and(pred == 1, gt == 0).sum().item()
290	        self.tn += torch.logical_and(pred == 0, gt == 0).sum().item()

322	        raise NotImplementedError("Unknown type {}".format(type(args)))
323	
324	
325	def stat_tensors(tensor):
326	    assert tensor.dim() == 1
327	    tensor_sm = tensor.softmax(0)
328	    entropy = (tensor_sm * torch.log(tensor_sm + 1e-9)).sum()
329	

506	        self.best_res = init_res
507	        self.best_ep = -1
508	
509	        self.better = better
510	        assert better in ["large", "small"]
511	
512	    def isbetter(self, new_res, old_res):
513	        if self.better == "large":

598	
599	def get_phrases_from_posmap(
600	    posmap: torch.BoolTensor, tokenized: Dict, tokenizer: AutoTokenizer
601	):
602	    assert isinstance(posmap, torch.Tensor), "posmap must be torch.Tensor"
603	    if posmap.dim() == 1:
604	        non_zero_idx = posmap.nonzero(as_tuple=True)[0].tolist()
605	        token_ids = [tokenized["input_ids"][i] for i in non_zero_idx]

75	            For large resolutions, 'binary_mask' may consume large amounts of
76	            memory.
77	        """
78	
79	        assert (points_per_side is None) != (
80	            point_grids is None
81	        ), "Exactly one of points_per_side or point_grid must be provided."
82	        if points_per_side is not None:
83	            self.point_grids = build_all_layer_point_grids(
84	                points_per_side,

89	            self.point_grids = point_grids
90	        else:
91	            raise ValueError("Can't have both points_per_side and point_grid be None.")
92	
93	        assert output_mode in [
94	            "binary_mask",
95	            "uncompressed_rle",
96	            "coco_rle",
97	        ], f"Unknown output_mode {output_mode}."
98	        if output_mode == "coco_rle":
99	            from pycocotools import mask as mask_utils  # type: ignore # noqa: F401
100	

215	                 resolution=(14, 14),
216	                 ):
217	        super().__init__()
218	        # (h, w)
219	        assert isinstance(resolution, tuple) and len(resolution) == 2
220	        self.num_heads = num_heads
221	        self.scale = key_dim ** -0.5
222	        self.key_dim = key_dim

306	        super().__init__()
307	        self.dim = dim
308	        self.input_resolution = input_resolution
309	        self.num_heads = num_heads
310	        assert window_size > 0, 'window_size must be greater than 0'
311	        self.window_size = window_size
312	        self.mlp_ratio = mlp_ratio
313	

313	
314	        self.drop_path = DropPath(
315	            drop_path) if drop_path > 0. else nn.Identity()
316	
317	        assert dim % num_heads == 0, 'dim must be divisible by num_heads'
318	        head_dim = dim // num_heads
319	
320	        window_resolution = (window_size, window_size)

332	
333	    def forward(self, x):
334	        H, W = self.input_resolution
335	        B, L, C = x.shape
336	        assert L == H * W, "input feature has wrong size"
337	        res_x = x
338	        if H == self.window_size and W == self.window_size:
339	            x = self.attn(x)

568	                i += 1
569	            if layer.downsample is not None:
570	                layer.downsample.apply(
571	                    lambda x: _set_lr_scale(x, lr_scales[i - 1]))
572	        assert i == depth
573	        for m in [self.norm_head, self.head]:
574	            m.apply(lambda x: _set_lr_scale(x, lr_scales[-1]))
575	

577	            p.param_name = k
578	
579	        def _check_lr_scale(m):
580	            for p in m.parameters():
581	                assert hasattr(p, 'lr_scale'), p.param_name
582	
583	        self.apply(_check_lr_scale)
584	

39	            1x3xHxW, which has been transformed with ResizeLongestSide.
40	          original_image_size (tuple(int, int)): The size of the image
41	            before transformation, in (H, W) format.
42	        """
43	        assert (
44	            len(transformed_image.shape) == 4
45	            and transformed_image.shape[1] == 3
46	            and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size
47	        ), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}."
48	        self.reset_image()
49	
50	        self.original_size = original_image_size