1	from ..utils import common_annotator_call, create_node_input_types, run_script
2	import comfy.model_management as model_management
3	import os, sys
4	import subprocess, threading
5	
6	def install_deps():
7	    try:
8	        import mediapipe

3	import numpy as np
4	import torch
5	from einops import rearrange
6	import os, sys
7	import subprocess, threading
8	import scipy.ndimage
9	import cv2
10	import torch.nn.functional as F

26	
27	    CATEGORY = "ControlNet Preprocessors/Optical Flow"
28	
29	    def estimate(self, image, ckpt_name, backward_flow=False, bidirectional_flow=False):
30	        assert len(image) > 1, "[Unimatch] Requiring as least two frames as a optical flow estimator. Only use this node on video input."    
31	        from controlnet_aux.unimatch import UnimatchDetector
32	        tensor_images = image
33	        model = UnimatchDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())

54	    CATEGORY = "ControlNet Preprocessors/Optical Flow"
55	    
56	    def mask_opt_flow(self, optical_flow, mask):
57	        from controlnet_aux.unimatch import flow_to_image
58	        assert len(mask) >= len(optical_flow), f"Not enough masks to mask optical flow: {len(mask)} vs {len(optical_flow)}"
59	        mask = mask[:optical_flow.shape[0]]
60	        mask = F.interpolate(mask, optical_flow.shape[1:3])
61	        mask = rearrange(mask, "n 1 h w -> n h w 1")

60	        original_type = type(box)
61	        is_numpy = isinstance(box, np.ndarray)
62	        single_box = isinstance(box, (list, tuple))
63	        if single_box:
64	            assert len(box) == 4 or len(box) == 5, (
65	                "BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor,"
66	                " where k == 4 or 5"
67	            )
68	            arr = torch.tensor(box)[None, :]
69	        else:
70	            # avoid modifying the input box

72	                arr = torch.from_numpy(np.asarray(box)).clone()
73	            else:
74	                arr = box.clone()
75	
76	        assert to_mode not in [BoxMode.XYXY_REL, BoxMode.XYWH_REL] and from_mode not in [
77	            BoxMode.XYXY_REL,
78	            BoxMode.XYWH_REL,
79	        ], "Relative mode not yet supported!"
80	
81	        if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:
82	            assert (

78	            BoxMode.XYWH_REL,
79	        ], "Relative mode not yet supported!"
80	
81	        if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:
82	            assert (
83	                arr.shape[-1] == 5
84	            ), "The last dimension of input shape must be 5 for XYWHA format"
85	            original_dtype = arr.dtype
86	            arr = arr.double()
87	

185	            matrix = cv2.resize(matrix, (w, h), self.interp_method_matrix)
186	        return mask, matrix
187	
188	    def _check_image(self, image_rgb):
189	        assert len(image_rgb.shape) == 3
190	        assert image_rgb.shape[2] == 3
191	        assert image_rgb.dtype == np.uint8
192	

186	        return mask, matrix
187	
188	    def _check_image(self, image_rgb):
189	        assert len(image_rgb.shape) == 3
190	        assert image_rgb.shape[2] == 3
191	        assert image_rgb.dtype == np.uint8
192	
193	    def _check_mask_matrix(self, mask, matrix):

187	
188	    def _check_image(self, image_rgb):
189	        assert len(image_rgb.shape) == 3
190	        assert image_rgb.shape[2] == 3
191	        assert image_rgb.dtype == np.uint8
192	
193	    def _check_mask_matrix(self, mask, matrix):
194	        assert len(matrix.shape) == 2

190	        assert image_rgb.shape[2] == 3
191	        assert image_rgb.dtype == np.uint8
192	
193	    def _check_mask_matrix(self, mask, matrix):
194	        assert len(matrix.shape) == 2
195	        assert len(mask.shape) == 2
196	        assert mask.dtype == np.uint8
197	

191	        assert image_rgb.dtype == np.uint8
192	
193	    def _check_mask_matrix(self, mask, matrix):
194	        assert len(matrix.shape) == 2
195	        assert len(mask.shape) == 2
196	        assert mask.dtype == np.uint8
197	
198	class DensePoseResultsVisualizer:

192	
193	    def _check_mask_matrix(self, mask, matrix):
194	        assert len(matrix.shape) == 2
195	        assert len(mask.shape) == 2
196	        assert mask.dtype == np.uint8
197	
198	class DensePoseResultsVisualizer:
199	    def visualize(

138	class DPT_DINOv2(nn.Module):
139	    def __init__(self, encoder='vitl', features=256, out_channels=[256, 512, 1024, 1024], use_bn=False, use_clstoken=False, localhub=True):
140	        super(DPT_DINOv2, self).__init__()
141	        
142	        assert encoder in ['vits', 'vitb', 'vitl']
143	        
144	        # in case the Internet connection is not stable, please load the DINOv2 locally
145	        if localhub:

26	    masks_list = []
27	    for i in range(0, n_samples_masked):
28	        prob_min = probs[i]
29	        prob_max = probs[i + 1]
30	        masks_list.append(torch.BoolTensor(mask_generator(int(N * random.uniform(prob_min, prob_max)))))
31	        upperbound += int(N * prob_max)
32	    for i in range(n_samples_masked, B):
33	        masks_list.append(torch.BoolTensor(mask_generator(0)))

44	            basename = f"ILSVRC2012_{self.value}_{actual_index:08d}"
45	        return os.path.join(dirname, basename + ".JPEG")
46	
47	    def parse_image_relpath(self, image_relpath: str) -> Tuple[str, int]:
48	        assert self != _Split.TEST
49	        dirname, filename = os.path.split(image_relpath)
50	        class_id = os.path.split(dirname)[-1]
51	        basename, _ = os.path.splitext(filename)

105	
106	    def _get_entries(self) -> np.ndarray:
107	        if self._entries is None:
108	            self._entries = self._load_extra(self._entries_path)
109	        assert self._entries is not None
110	        return self._entries
111	
112	    def _get_class_ids(self) -> np.ndarray:

110	        return self._entries
111	
112	    def _get_class_ids(self) -> np.ndarray:
113	        if self._split == _Split.TEST:
114	            assert False, "Class IDs are not available in TEST split"
115	        if self._class_ids is None:
116	            self._class_ids = self._load_extra(self._class_ids_path)
117	        assert self._class_ids is not None

113	        if self._split == _Split.TEST:
114	            assert False, "Class IDs are not available in TEST split"
115	        if self._class_ids is None:
116	            self._class_ids = self._load_extra(self._class_ids_path)
117	        assert self._class_ids is not None
118	        return self._class_ids
119	
120	    def _get_class_names(self) -> np.ndarray:

118	        return self._class_ids
119	
120	    def _get_class_names(self) -> np.ndarray:
121	        if self._split == _Split.TEST:
122	            assert False, "Class names are not available in TEST split"
123	        if self._class_names is None:
124	            self._class_names = self._load_extra(self._class_names_path)
125	        assert self._class_names is not None

121	        if self._split == _Split.TEST:
122	            assert False, "Class names are not available in TEST split"
123	        if self._class_names is None:
124	            self._class_names = self._load_extra(self._class_names_path)
125	        assert self._class_names is not None
126	        return self._class_names
127	
128	    def find_class_id(self, class_index: int) -> str:

165	        return None if self.split == _Split.TEST else str(class_name)
166	
167	    def __len__(self) -> int:
168	        entries = self._get_entries()
169	        assert len(entries) == self.split.length
170	        return len(entries)
171	
172	    def _load_labels(self, labels_path: str) -> List[Tuple[str, str]]:

232	                entries_array[index] = (actual_index, class_index, class_id, class_name)
233	        else:
234	            class_names = {class_id: class_name for class_id, class_name in labels}
235	
236	            assert dataset
237	            old_percent = -1
238	            for index in range(sample_count):
239	                percent = 100 * (index + 1) // sample_count

166	                        class_entries.append(class_entry)
167	            except OSError as e:
168	                raise RuntimeError(f'can not read blocks file "{path}"') from e
169	
170	            assert class_entries[-1].maybe_filename is None
171	
172	            for class_entry1, class_entry2 in zip(class_entries, class_entries[1:]):
173	                assert class_entry1.block_offset <= class_entry2.block_offset

169	
170	            assert class_entries[-1].maybe_filename is None
171	
172	            for class_entry1, class_entry2 in zip(class_entries, class_entries[1:]):
173	                assert class_entry1.block_offset <= class_entry2.block_offset
174	                start_offset = 512 * class_entry1.block_offset
175	                end_offset = 512 * class_entry2.block_offset
176	                assert class_entry1.maybe_filename is not None

172	            for class_entry1, class_entry2 in zip(class_entries, class_entries[1:]):
173	                assert class_entry1.block_offset <= class_entry2.block_offset
174	                start_offset = 512 * class_entry1.block_offset
175	                end_offset = 512 * class_entry2.block_offset
176	                assert class_entry1.maybe_filename is not None
177	                filename = class_entry1.maybe_filename
178	                entry = _Entry(class_index, start_offset, end_offset, filename)
179	                # Skip invalid image files (PIL throws UnidentifiedImageError)

211	            mapped_data = class_mmap[start_offset:end_offset]
212	            data = mapped_data[512:]  # Skip entry header block
213	
214	            if len(data) >= 2 and tuple(data[:2]) == (0x1F, 0x8B):
215	                assert index in self._gzipped_indices, f"unexpected gzip header for sample {index}"
216	                with GzipFile(fileobj=BytesIO(data)) as g:
217	                    data = g.read()
218	        except Exception as e:

49	    kwargs = {}
50	
51	    for token in tokens[1:]:
52	        key, value = token.split("=")
53	        assert key in ("root", "extra", "split")
54	        kwargs[key] = value
55	
56	    if name == "ImageNet":

49	
50	    def _mask(self, mask, max_mask_patches):
51	        delta = 0
52	        for _ in range(10):
53	            target_area = random.uniform(self.min_num_patches, max_mask_patches)
54	            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
55	            h = int(round(math.sqrt(target_area * aspect_ratio)))
56	            w = int(round(math.sqrt(target_area / aspect_ratio)))

50	    def _mask(self, mask, max_mask_patches):
51	        delta = 0
52	        for _ in range(10):
53	            target_area = random.uniform(self.min_num_patches, max_mask_patches)
54	            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
55	            h = int(round(math.sqrt(target_area * aspect_ratio)))
56	            w = int(round(math.sqrt(target_area / aspect_ratio)))
57	            if w < self.width and h < self.height:

54	            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
55	            h = int(round(math.sqrt(target_area * aspect_ratio)))
56	            w = int(round(math.sqrt(target_area / aspect_ratio)))
57	            if w < self.width and h < self.height:
58	                top = random.randint(0, self.height - h)
59	                left = random.randint(0, self.width - w)
60	
61	                num_masked = mask[top : top + h, left : left + w].sum()

55	            h = int(round(math.sqrt(target_area * aspect_ratio)))
56	            w = int(round(math.sqrt(target_area / aspect_ratio)))
57	            if w < self.width and h < self.height:
58	                top = random.randint(0, self.height - h)
59	                left = random.randint(0, self.width - w)
60	
61	                num_masked = mask[top : top + h, left : left + w].sum()
62	                # Overlap

102	
103	        yield from itertools.islice(iterator, self._advance, None)
104	
105	    def _iterator(self):
106	        assert not self._shuffle
107	
108	        while True:
109	            iterable = range(self._sample_count)

109	            iterable = range(self._sample_count)
110	            yield from itertools.islice(iterable, self._start, None, self._step)
111	
112	    def _shuffled_iterator(self):
113	        assert self._shuffle
114	
115	        # Instantiate a generator here (rather than in the ctor) to keep the class
116	        # picklable (requirement of mp.spawn)

199	
200	        yield from itertools.islice(iterator, self._advance, None)
201	
202	    def _iterator(self):
203	        assert not self._shuffle
204	
205	        while True:
206	            iterable = range(self._sample_count)

206	            iterable = range(self._sample_count)
207	            yield from itertools.islice(iterable, self._start, None, self._step)
208	
209	    def _shuffled_iterator(self):
210	        assert self._shuffle
211	
212	        # Instantiate a generator here (rather than in the ctor) to be keep the class
213	        # picklable (requirement of mp.spawn)

47	        The rank of the current process within the local (per-machine) process group.
48	    """
49	    if not is_enabled():
50	        return 0
51	    assert 0 <= _LOCAL_RANK < _LOCAL_WORLD_SIZE
52	    return _LOCAL_RANK
53	
54	

59	        i.e. the number of processes per machine.
60	    """
61	    if not is_enabled():
62	        return 1
63	    assert 0 <= _LOCAL_RANK < _LOCAL_WORLD_SIZE
64	    return _LOCAL_WORLD_SIZE
65	
66	

92	    MIN_MASTER_PORT, MAX_MASTER_PORT = (20_000, 60_000)
93	
94	    master_port_str = os.environ.get("MASTER_PORT")
95	    if master_port_str is None:
96	        rng = random.Random(seed)
97	        return rng.randint(MIN_MASTER_PORT, MAX_MASTER_PORT)
98	
99	    return int(master_port_str)

183	        # logger.info("Initialization from Slurm environment")
184	        job_id = int(os.environ["SLURM_JOB_ID"])
185	        node_count = int(os.environ["SLURM_JOB_NUM_NODES"])
186	        nodes = _parse_slurm_node_list(os.environ["SLURM_JOB_NODELIST"])
187	        assert len(nodes) == node_count
188	
189	        self.master_addr = nodes[0]
190	        self.master_port = _get_master_port(seed=job_id)

189	        self.master_addr = nodes[0]
190	        self.master_port = _get_master_port(seed=job_id)
191	        self.rank = int(os.environ["SLURM_PROCID"])
192	        self.world_size = int(os.environ["SLURM_NTASKS"])
193	        assert self.rank < self.world_size
194	        self.local_rank = int(os.environ["SLURM_LOCALID"])
195	        self.local_world_size = self.world_size // node_count
196	        assert self.local_rank < self.local_world_size

192	        self.world_size = int(os.environ["SLURM_NTASKS"])
193	        assert self.rank < self.world_size
194	        self.local_rank = int(os.environ["SLURM_LOCALID"])
195	        self.local_world_size = self.world_size // node_count
196	        assert self.local_rank < self.local_world_size
197	
198	    # Single node job with preset environment (i.e. torchrun)
199	    def _set_from_preset_env(self):

201	        self.master_addr = os.environ["MASTER_ADDR"]
202	        self.master_port = os.environ["MASTER_PORT"]
203	        self.rank = int(os.environ["RANK"])
204	        self.world_size = int(os.environ["WORLD_SIZE"])
205	        assert self.rank < self.world_size
206	        self.local_rank = int(os.environ["LOCAL_RANK"])
207	        self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
208	        assert self.local_rank < self.local_world_size

204	        self.world_size = int(os.environ["WORLD_SIZE"])
205	        assert self.rank < self.world_size
206	        self.local_rank = int(os.environ["LOCAL_RANK"])
207	        self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
208	        assert self.local_rank < self.local_world_size
209	
210	    # Single node and GPU job (i.e. local script run)
211	    def _set_from_local(self):

171	    def forward(self, features_rank):
172	        """
173	        Compute the results on all values of `self.nb_knn` neighbors from the full `self.max_k`
174	        """
175	        assert all(k <= self.max_k for k in self.nb_knn)
176	
177	        topk_sims, neighbors_labels = self.compute_neighbors(features_rank)
178	        batch_size = neighbors_labels.shape[0]

488	        test_dataset_strs = [val_dataset_str]
489	    if test_metric_types is None:
490	        test_metric_types = [val_metric_type] * len(test_dataset_strs)
491	    else:
492	        assert len(test_metric_types) == len(test_dataset_strs)
493	    assert len(test_dataset_strs) == len(test_class_mapping_fpaths)
494	
495	    train_transform = make_classification_train_transform()

489	    if test_metric_types is None:
490	        test_metric_types = [val_metric_type] * len(test_dataset_strs)
491	    else:
492	        assert len(test_metric_types) == len(test_dataset_strs)
493	    assert len(test_dataset_strs) == len(test_class_mapping_fpaths)
494	
495	    train_transform = make_classification_train_transform()
496	    train_dataset = make_dataset(

141	
142	    logger.info(f"Features shape: {tuple(features.shape)}")
143	    logger.info(f"Labels shape: {tuple(all_labels.shape)}")
144	
145	    assert torch.all(all_labels > -1)
146	
147	    return features, all_labels

105	        data.update(kwargs)
106	
107	        basename = f"{name}.{rankstr()}.pth"
108	        save_file = os.path.join(self.save_dir, basename)
109	        assert os.path.basename(save_file) == basename, basename
110	        self.logger.info("Saving checkpoint to {}".format(save_file))
111	        with self.path_manager.open(save_file, "wb") as f:
112	            torch.save(data, f)

58	
59	class MemEffAttention(Attention):
60	    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
61	        if not XFORMERS_AVAILABLE:
62	            assert attn_bias is None, "xFormers is required for nested tensors usage"
63	            return super().forward(x)
64	
65	        B, N, C = x.shape

198	    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
199	        """
200	        x_list contains a list of tensors to nest together and run
201	        """
202	        assert isinstance(self.attn, MemEffAttention)
203	
204	        if self.training and self.sample_drop_ratio > 0.0:
205	

238	    def forward(self, x_or_x_list):
239	        if isinstance(x_or_x_list, Tensor):
240	            return super().forward(x_or_x_list)
241	        elif isinstance(x_or_x_list, list):
242	            assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
243	            return self.forward_nested(x_or_x_list)
244	        else:
245	            raise AssertionError

15	
16	
17	def make_2tuple(x):
18	    if isinstance(x, tuple):
19	        assert len(x) == 2
20	        return x
21	
22	    assert isinstance(x, int)

18	    if isinstance(x, tuple):
19	        assert len(x) == 2
20	        return x
21	
22	    assert isinstance(x, int)
23	    return (x, x)
24	
25	

69	    def forward(self, x: Tensor) -> Tensor:
70	        _, _, H, W = x.shape
71	        patch_H, patch_W = self.patch_size
72	
73	        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
74	        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
75	
76	        x = self.proj(x)  # B C H W

70	        _, _, H, W = x.shape
71	        patch_H, patch_W = self.patch_size
72	
73	        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
74	        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
75	
76	        x = self.proj(x)  # B C H W
77	        H, W = x.size(2), x.size(3)

27	    def update(self, **kwargs):
28	        for k, v in kwargs.items():
29	            if isinstance(v, torch.Tensor):
30	                v = v.item()
31	            assert isinstance(v, (float, int))
32	            self.meters[k].update(v)
33	
34	    def __getattr__(self, attr):

183	            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
184	            mode="bicubic",
185	        )
186	
187	        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
188	        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
189	        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
190	

243	        for i, blk in enumerate(self.blocks):
244	            x = blk(x)
245	            if i in blocks_to_take:
246	                output.append(x)
247	        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
248	        return output
249	
250	    def _get_intermediate_layers_chunked(self, x, n=1):

257	                x = blk(x)
258	                if i in blocks_to_take:
259	                    output.append(x)
260	                i += 1
261	        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
262	        return output
263	
264	    def get_intermediate_layers(

50	    args = args_parser.parse_args()
51	
52	    setup_logging()
53	
54	    assert os.path.exists(args.config_file), "Configuration file does not exist!"
55	    submit_jobs(Evaluator, args, name="dinov2:knn")
56	    return 0
57	

50	    args = args_parser.parse_args()
51	
52	    setup_logging()
53	
54	    assert os.path.exists(args.config_file), "Configuration file does not exist!"
55	    submit_jobs(Evaluator, args, name="dinov2:linear")
56	    return 0
57	

50	    args = args_parser.parse_args()
51	
52	    setup_logging()
53	
54	    assert os.path.exists(args.config_file), "Configuration file does not exist!"
55	    submit_jobs(Evaluator, args, name="dinov2:logreg")
56	    return 0
57	

50	    args = args_parser.parse_args()
51	
52	    setup_logging()
53	
54	    assert os.path.exists(args.config_file), "Configuration file does not exist!"
55	    submit_jobs(Trainer, args, name="dinov2:train")
56	    return 0
57	

24	
25	    XFORMERS_AVAILABLE = True
26	except ImportError:
27	    XFORMERS_AVAILABLE = False
28	assert XFORMERS_AVAILABLE, "xFormers is required for DINOv2 training"
29	
30	
31	logger = logging.getLogger("dinov2")

90	        logger.info(f"OPTIONS -- IBOT masking -- ibot_mask_ratio_tuple: {cfg.ibot.mask_ratio_min_max}")
91	        logger.info(f"OPTIONS -- IBOT masking -- ibot_mask_sample_probability: {cfg.ibot.mask_sample_probability}")
92	        if self.do_ibot:
93	            self.ibot_loss_weight = cfg.ibot.loss_weight
94	            assert max(cfg.ibot.mask_ratio_min_max) > 0, "please provide a positive mask ratio tuple for ibot"
95	            assert cfg.ibot.mask_sample_probability > 0, "please provide a positive mask probability for ibot"
96	            self.ibot_out_dim = cfg.ibot.head_n_prototypes if self.ibot_separate_head else cfg.dino.head_n_prototypes
97	            self.ibot_patch_loss = iBOTPatchLoss(self.ibot_out_dim)

91	        logger.info(f"OPTIONS -- IBOT masking -- ibot_mask_sample_probability: {cfg.ibot.mask_sample_probability}")
92	        if self.do_ibot:
93	            self.ibot_loss_weight = cfg.ibot.loss_weight
94	            assert max(cfg.ibot.mask_ratio_min_max) > 0, "please provide a positive mask ratio tuple for ibot"
95	            assert cfg.ibot.mask_sample_probability > 0, "please provide a positive mask probability for ibot"
96	            self.ibot_out_dim = cfg.ibot.head_n_prototypes if self.ibot_separate_head else cfg.dino.head_n_prototypes
97	            self.ibot_patch_loss = iBOTPatchLoss(self.ibot_out_dim)
98	            if self.ibot_separate_head:

133	            loss.backward()
134	
135	    def forward_backward(self, images, teacher_temp):
136	        n_global_crops = 2
137	        assert n_global_crops == 2
138	        n_local_crops = self.cfg.crops.local_crops_number
139	
140	        global_crops = images["collated_global_crops"].cuda(non_blocking=True)

54	    if checkpoint_path is None:
55	        return None
56	
57	    username = os.environ.get("USER")
58	    assert username is not None
59	    return checkpoint_path / username
60	
61	

33	    if isinstance(dtype, torch.dtype):
34	        return dtype
35	    if isinstance(dtype, str):
36	        dtype = np.dtype(dtype)
37	    assert isinstance(dtype, np.dtype), f"Expected an instance of nunpy dtype, got {type(dtype)}"
38	    return _NUMPY_TO_TORCH_DTYPE[dtype]

6	
7	import logging
8	import os
9	import random
10	import subprocess
11	from urllib.parse import urlparse
12	
13	import numpy as np

47	def get_sha():
48	    cwd = os.path.dirname(os.path.abspath(__file__))
49	
50	    def _run(command):
51	        return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
52	
53	    sha = "N/A"
54	    diff = "clean"

54	    diff = "clean"
55	    branch = "N/A"
56	    try:
57	        sha = _run(["git", "rev-parse", "HEAD"])
58	        subprocess.check_output(["git", "diff"], cwd=cwd)
59	        diff = _run(["git", "diff-index", "HEAD"])
60	        diff = "has uncommitted changes" if diff else "clean"
61	        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])

54	    diff = "clean"
55	    branch = "N/A"
56	    try:
57	        sha = _run(["git", "rev-parse", "HEAD"])
58	        subprocess.check_output(["git", "diff"], cwd=cwd)
59	        diff = _run(["git", "diff-index", "HEAD"])
60	        diff = "has uncommitted changes" if diff else "clean"
61	        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])

58	        subprocess.check_output(["git", "diff"], cwd=cwd)
59	        diff = _run(["git", "diff-index", "HEAD"])
60	        diff = "has uncommitted changes" if diff else "clean"
61	        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
62	    except Exception:
63	        pass
64	    message = f"sha: {sha}, status: {diff}, branch: {branch}"
65	    return message
66	

78	        iters = np.arange(total_iters - warmup_iters - freeze_iters)
79	        schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))
80	        self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule))
81	
82	        assert len(self.schedule) == self.total_iters
83	
84	    def __getitem__(self, it):
85	        if it >= self.total_iters:

107	        num_patches = self.patch_embed.num_patches
108	
109	        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
110	        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
111	        assert num_register_tokens >= 0
112	        self.register_tokens = (
113	            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
114	        )

203	            mode="bicubic",
204	            antialias=self.interpolate_antialias
205	        )
206	        
207	        assert int(w0) == patch_pos_embed.shape[-2]
208	        assert int(h0) == patch_pos_embed.shape[-1]
209	        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
210	        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)

204	            antialias=self.interpolate_antialias
205	        )
206	        
207	        assert int(w0) == patch_pos_embed.shape[-2]
208	        assert int(h0) == patch_pos_embed.shape[-1]
209	        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
210	        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
211	

276	        for i, blk in enumerate(self.blocks):
277	            x = blk(x)
278	            if i in blocks_to_take:
279	                output.append(x)
280	        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
281	        return output
282	
283	    def _get_intermediate_layers_chunked(self, x, n=1):

290	                x = blk(x)
291	                if i in blocks_to_take:
292	                    output.append(x)
293	                i += 1
294	        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
295	        return output
296	
297	    def get_intermediate_layers(

12	import torchvision.transforms as transforms
13	import torchvision.datasets as datasets
14	from typing import Any, Callable, Optional, Tuple
15	import os
16	import pickle
17	import numpy as np
18	import copy
19	import custom_albumentations as albumentations

111	        # print("lb_data.size:", lb_data.size)
112	        lb = np.array(lb_data, dtype=np.float32)
113	        if lb.ndim == 3:
114	            lb = np.squeeze(lb[:, :, 0])
115	        assert lb.ndim == 2
116	        threshold = self.threshold
117	        lb = lb[np.newaxis, :, :]
118	

86	        self.image_size = image_size
87	        self.train_sample = train_sample
88	        self.objective = objective
89	        self.start_dist = start_dist
90	        assert start_dist in ['normal', 'uniform']
91	
92	        assert objective in {'pred_noise', 'pred_x0', 'pred_v', 'pred_delta', 'pred_KC'}, 'objective must be either pred_noise (predict noise) or pred_x0 (predict image start) or pred_v (predict v [v-parameterization as defined in appendix D of progressive distillation paper, used in imagen-video successfully])'
93	

88	        self.objective = objective
89	        self.start_dist = start_dist
90	        assert start_dist in ['normal', 'uniform']
91	
92	        assert objective in {'pred_noise', 'pred_x0', 'pred_v', 'pred_delta', 'pred_KC'}, 'objective must be either pred_noise (predict noise) or pred_x0 (predict image start) or pred_v (predict v [v-parameterization as defined in appendix D of progressive distillation paper, used in imagen-video successfully])'
93	
94	        if beta_schedule == 'linear':
95	            betas = linear_beta_schedule(timesteps)

148	
149	        # calculate p2 reweighting
150	
151	        register_buffer('p2_loss_weight', (p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod)) ** -p2_loss_weight_gamma)
152	        assert not torch.isnan(self.p2_loss_weight).all()
153	        if self.objective == "pred_noise":
154	            lvlb_weights = self.betas ** 2 / (
155	                        2 * (self.posterior_variance+1e-5) * alphas * (1 - self.alphas_cumprod))

161	            lvlb_weights = 0.5 * torch.sqrt(alphas_cumprod) / (2. * 1 - alphas_cumprod)
162	        elif self.objective == "pred_v":
163	            lvlb_weights = 0.5 * torch.sqrt(alphas_cumprod) / (2. * 1 - alphas_cumprod)
164	        self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
165	        assert not torch.isnan(self.lvlb_weights).all()
166	        self.use_l1 = use_l1
167	
168	        self.perceptual_weight = perceptual_weight

266	    def interpolate(self, x1, x2, mask, t = None, lam = 0.5):
267	        b, *_, device = *x1.shape, x1.device
268	        t = default(t, self.num_timesteps - 1)
269	
270	        assert x1.shape == x2.shape
271	
272	        t_batched = torch.stack([torch.tensor(t, device = device)] * b)
273	        xt1, xt2 = map(lambda x: self.q_sample(x, t = t_batched), (x1, x2))

277	            img = self.p_sample(img, mask, torch.full((b,), i, device=device, dtype=torch.long))
278	        return img
279	
280	    def get_input(self, batch, return_first_stage_outputs=False, return_original_cond=False):
281	        assert 'image' in self.input_keys;
282	        if len(self.input_keys) > len(batch.keys()):
283	            x, *_ = batch.values()
284	        else:

559	        ckpt_path = kwargs.pop("ckpt_path", None)
560	        ignore_keys = kwargs.pop("ignore_keys", [])
561	        only_model = kwargs.pop("only_model", False)
562	        super().__init__(*args, **kwargs)
563	        assert self.num_timesteps_cond <= self.num_timesteps
564	        if not scale_by_std:
565	            self.scale_factor = scale_factor
566	        else:

567	            self.register_buffer('scale_factor', torch.tensor(scale_factor))
568	        if self.scale_by_softsign:
569	            self.scale_by_std = False
570	            print('### USING SOFTSIGN RESCALING')
571	        assert (self.scale_by_std and self.scale_by_softsign) is False;
572	
573	        self.init_first_stage(auto_encoder)
574	        # self.instantiate_cond_stage(cond_stage_config)

573	        self.init_first_stage(auto_encoder)
574	        # self.instantiate_cond_stage(cond_stage_config)
575	        self.input_keys = input_keys
576	        self.clip_denoised = False
577	        assert sample_type in ['p_loop', 'ddim', 'dpm', 'transformer'] ###  'dpm' is not availible now, suggestion 'ddim'
578	        self.sample_type = sample_type
579	
580	        if ckpt_path is not None:

620	    def on_train_batch_start(self, batch):
621	        # only for the first batch
622	        if self.scale_by_std and (not self.scale_by_softsign):
623	            if not self.default_scale:
624	                assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
625	                # set rescale weight to 1./std of encodings
626	                print("### USING STD-RESCALING ###")
627	                x, *_ = batch.values()

637	            print(f'### USING SOFTSIGN SCALE !')
638	
639	    @torch.no_grad()
640	    def get_input(self, batch, return_first_stage_outputs=False, return_original_cond=False):
641	        assert 'image' in self.input_keys;
642	        # if len(self.input_keys) > len(batch.keys()):
643	        #     x, cond, *_ = batch.values()
644	        # else:

7	    return val is not None
8	
9	
10	def clamp(value, min_value=None, max_value=None):
11	    assert exists(min_value) or exists(max_value)
12	    if exists(min_value):
13	        value = max(value, min_value)
14	

90	        self.inv_gamma = inv_gamma
91	        self.power = power
92	        self.min_value = min_value
93	
94	        assert isinstance(param_or_buffer_names_no_ema, (set, list))
95	        self.param_or_buffer_names_no_ema = param_or_buffer_names_no_ema  # parameter or buffer
96	
97	        self.ignore_names = ignore_names

34	    Build sinusoidal embeddings.
35	    This matches the implementation in tensor2tensor, but differs slightly
36	    from the description in Section 3.5 of "Attention Is All You Need".
37	    """
38	    assert len(timesteps.shape) == 1
39	
40	    half_dim = embedding_dim // 2
41	    emb = math.log(10000) / (half_dim - 1)

220	        return x+h_
221	
222	
223	def make_attn(in_channels, attn_type="vanilla"):
224	    assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
225	    print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
226	    if attn_type == "vanilla":
227	        return AttnBlock(in_channels)

337	            # assume aligned context, cat along channel axis
338	            x = torch.cat((x, context), dim=1)
339	        if self.use_timestep:
340	            # timestep embedding
341	            assert t is not None
342	            temb = get_timestep_embedding(t, self.ch)
343	            temb = self.temb.dense[0](temb)
344	            temb = nonlinearity(temb)

745	
746	class Upsampler(nn.Module):
747	    def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
748	        super().__init__()
749	        assert out_size >= in_size
750	        num_blocks = int(np.log2(out_size//in_size))+1
751	        factor_up = 1.+ (out_size % in_size)
752	        print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")

769	        self.mode = mode
770	        if self.with_conv:
771	            print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
772	            raise NotImplementedError()
773	            assert in_channels is not None
774	            # no asymmetric padding in torch conv, must do it ourselves
775	            self.conv = torch.nn.Conv2d(in_channels,
776	                                        in_channels,

794	                 dropout=0.,
795	                 pretrained_config=None):
796	        super().__init__()
797	        if pretrained_config is None:
798	            assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
799	            self.pretrained_model = pretrained_model
800	        else:
801	            assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'

797	        if pretrained_config is None:
798	            assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
799	            self.pretrained_model = pretrained_model
800	        else:
801	            assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
802	            self.instantiate_pretrained(pretrained_config)
803	
804	        self.do_reshape = reshape

907	        self.encoder = Encoder(**ddconfig)
908	        self.decoder = Decoder(**ddconfig)
909	        self.down_ratio = 2 ** (len(ddconfig['ch_mult']) - 1)
910	        self.loss = LPIPSWithDiscriminator(**lossconfig)
911	        assert ddconfig["double_z"]
912	        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
913	        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
914	        self.embed_dim = embed_dim

912	        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
913	        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
914	        self.embed_dim = embed_dim
915	        if colorize_nlabels is not None:
916	            assert type(colorize_nlabels)==int
917	            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
918	        if monitor is not None:
919	            self.monitor = monitor

1	import os, yaml, pickle, shutil, tarfile, glob
2	import cv2
3	import custom_albumentations as albumentations
4	import PIL
5	import numpy as np
6	import torchvision.transforms.functional as TF
7	# from omegaconf import OmegaConf
8	from functools import partial

83	        if (not os.path.exists(self.human2integer)):
84	            download(URL, self.human2integer)
85	        with open(self.human2integer, "r") as f:
86	            lines = f.read().splitlines()
87	            assert len(lines) == 1000
88	            self.human2integer_dict = dict()
89	            for line in lines:
90	                value, key = line.split(":")

169	                path = os.path.join(self.root, self.FILES[0])
170	                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
171	                    import academictorrents as at
172	                    atpath = at.get(self.AT_HASH, datastore=self.root)
173	                    assert atpath == path
174	
175	                print("Extracting {} to {}".format(path, datadir))
176	                os.makedirs(datadir, exist_ok=True)

234	                path = os.path.join(self.root, self.FILES[0])
235	                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
236	                    import academictorrents as at
237	                    atpath = at.get(self.AT_HASH, datastore=self.root)
238	                    assert atpath == path
239	
240	                print("Extracting {} to {}".format(path, datadir))
241	                os.makedirs(datadir, exist_ok=True)

290	        :param data_root:
291	        :param random_crop:
292	        """
293	        self.base = self.get_base()
294	        assert size
295	        assert (size / downscale_f).is_integer()
296	        self.size = size
297	        self.LR_size = int(size / downscale_f)

291	        :param random_crop:
292	        """
293	        self.base = self.get_base()
294	        assert size
295	        assert (size / downscale_f).is_integer()
296	        self.size = size
297	        self.LR_size = int(size / downscale_f)
298	        self.min_crop_f = min_crop_f

296	        self.size = size
297	        self.LR_size = int(size / downscale_f)
298	        self.min_crop_f = min_crop_f
299	        self.max_crop_f = max_crop_f
300	        assert(max_crop_f <= 1.)
301	        self.center_crop = not random_crop
302	
303	        self.image_rescaler = albumentations.SmallestMaxSize(max_size=size, interpolation=cv2.INTER_AREA)

11	                 perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
12	                 disc_loss="hinge"):
13	
14	        super().__init__()
15	        assert disc_loss in ["hinge", "vanilla"]
16	        self.kl_weight = kl_weight
17	        self.pixel_weight = pixelloss_weight
18	        self.perceptual_loss = LPIPS().eval()

64	        # now the GAN part
65	        if optimizer_idx == 0:
66	            # generator update
67	            if cond is None:
68	                assert not self.disc_conditional
69	                logits_fake = self.discriminator(reconstructions.contiguous())
70	            else:
71	                assert self.disc_conditional

67	            if cond is None:
68	                assert not self.disc_conditional
69	                logits_fake = self.discriminator(reconstructions.contiguous())
70	            else:
71	                assert self.disc_conditional
72	                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1))
73	            g_loss = -torch.mean(logits_fake)
74	

75	            if self.disc_factor > 0.0:
76	                try:
77	                    d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
78	                except RuntimeError:
79	                    assert not self.training
80	                    d_weight = torch.tensor(0.0)
81	            else:
82	                d_weight = torch.tensor(0.0)

41	        num_pos_feats = num_feats // 2
42	        # mask = tensor_list.mask
43	        mask = torch.zeros(x.shape[0], x.shape[1], x.shape[2], device=x.device).to(torch.bool)
44	        batch = mask.shape[0]
45	        assert mask is not None
46	        not_mask = ~mask
47	        y_embed = not_mask.cumsum(1, dtype=torch.float32)
48	        x_embed = not_mask.cumsum(2, dtype=torch.float32)

414	    """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """
415	
416	    def __init__(self, dim, is_random = False):
417	        super().__init__()
418	        assert (dim % 2) == 0
419	        half_dim = dim // 2
420	        self.weights = nn.Parameter(torch.randn(half_dim), requires_grad = not is_random)
421	

35	
36	class weight_quantize_fn(nn.Module):
37	    def __init__(self, bit_w):
38	        super(weight_quantize_fn, self).__init__()
39	        assert bit_w > 0
40	
41	        self.bit_w = bit_w - 1
42	        self.weight_q = weight_quantization(b=self.bit_w)

87	        self.bit_a = bit_a
88	        self.signed = signed
89	        if signed:
90	            self.bit_a -= 1
91	        assert bit_a > 0
92	
93	        self.act_q = act_quantization(b=self.bit_a, signed=signed)
94	        self.register_parameter('a_alpha', Parameter(torch.tensor(8.0), requires_grad=True))

122	    """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """
123	
124	    def __init__(self, dim, is_random = False):
125	        super().__init__()
126	        assert (dim % 2) == 0
127	        half_dim = dim // 2
128	        self.weights = nn.Parameter(torch.randn(half_dim), requires_grad = not is_random)
129	

191	                           (int(x1), int(preds.shape[3] - x2), int(y1),
192	                            int(preds.shape[2] - y2)))
193	            count_mat[:, :, y1:y2, x1:x2] += 1
194	
195	        assert (count_mat == 0).sum() == 0
196	        seg_logits = preds / count_mat
197	        return seg_logits

43	            self.img_id_to_captions = dict()
44	            self.img_id_to_filepath = dict()
45	            self.img_id_to_segmentation_filepath = dict()
46	
47	        assert data_json.split("/")[-1] in ["captions_train2017.json",
48	                                            "captions_val2017.json"]
49	        if self.stuffthing:
50	            self.segmentation_prefix = (
51	                "data/cocostuffthings/val2017" if

103	        if not self.onehot and not segmentation.mode == "RGB":
104	            segmentation = segmentation.convert("RGB")
105	        segmentation = np.array(segmentation).astype(np.uint8)
106	        if self.onehot:
107	            assert self.stuffthing
108	            # stored in caffe format: unlabeled==255. stuff and thing from
109	            # 0-181. to be compatible with the labels in
110	            # https://github.com/nightrome/cocostuff/blob/master/labels.txt

110	            # https://github.com/nightrome/cocostuff/blob/master/labels.txt
111	            # we shift stuffthing one to the right and put unlabeled in zero
112	            # as long as segmentation is uint8 shifting to right handles the
113	            # latter too
114	            assert segmentation.dtype == np.uint8
115	            segmentation = segmentation + 1
116	
117	        processed = self.preprocessor(image=image, segmentation=segmentation)

118	        image, segmentation = processed["image"], processed["segmentation"]
119	        image = (image / 127.5 - 1.0).astype(np.float32)
120	
121	        if self.onehot:
122	            assert segmentation.dtype == np.uint8
123	            # make it one hot
124	            n_labels = 183
125	            flatseg = np.ravel(segmentation)

32	        if self.encode_crop:
33	            crop_coordinates = self.bbox_from_token_pair(conditional_list[-2], conditional_list[-1])
34	            conditional_list = conditional_list[:-2]
35	        object_triples = grouper(conditional_list, 3)
36	        assert conditional.shape[0] == self.embedding_dim
37	        return [
38	            (object_triple[0], self.bbox_from_token_pair(object_triple[1], object_triple[2]))
39	            for object_triple in object_triples if object_triple[0] != self.none

76	        if self.encode_crop:
77	            crop_coordinates = self.bbox_from_token_pair(conditional_list[-2], conditional_list[-1])
78	            conditional_list = conditional_list[:-2]
79	        table_of_content = grouper(conditional_list, self.object_descriptor_length)
80	        assert conditional.shape[0] == self.embedding_dim
81	        return [
82	            (object_tuple[0], self.coordinates_from_token(object_tuple[1]))
83	            for object_tuple in table_of_content if object_tuple[0] != self.none

162	            extra = []
163	
164	        object_tuples = self._make_object_descriptors(annotations)
165	        flattened = [token for tuple_ in object_tuples for token in tuple_] + extra
166	        assert len(flattened) == self.embedding_dim
167	        assert all(0 <= value < self.no_tokens for value in flattened)
168	        return LongTensor(flattened)

163	
164	        object_tuples = self._make_object_descriptors(annotations)
165	        flattened = [token for tuple_ in object_tuples for token in tuple_] + extra
166	        assert len(flattened) == self.embedding_dim
167	        assert all(0 <= value < self.no_tokens for value in flattened)
168	        return LongTensor(flattened)

74	        max_size = min(width, height)
75	        if max_size <= self.min_size:
76	            size = max_size
77	        else:
78	            size = random.randint(self.min_size, max_size)
79	        top = random.randint(0, height - size)
80	        left = random.randint(0, width - size)
81	        bbox = left / width, top / height, size / width, size / height

75	        if max_size <= self.min_size:
76	            size = max_size
77	        else:
78	            size = random.randint(self.min_size, max_size)
79	        top = random.randint(0, height - size)
80	        left = random.randint(0, width - size)
81	        bbox = left / width, top / height, size / width, size / height
82	        return bbox, F.crop(img, top, left, size, size)

76	            size = max_size
77	        else:
78	            size = random.randint(self.min_size, max_size)
79	        top = random.randint(0, height - size)
80	        left = random.randint(0, width - size)
81	        bbox = left / width, top / height, size / width, size / height
82	        return bbox, F.crop(img, top, left, size, size)
83	

23	
24	
25	def str_to_indices(string):
26	    """Expects a string in the format '32-123, 256, 280-321'"""
27	    assert not string.endswith(","), "provided string '{}' ends with a comma, pls remove it".format(string)
28	    subs = string.split(",")
29	    indices = []
30	    for sub in subs:

28	    subs = string.split(",")
29	    indices = []
30	    for sub in subs:
31	        subsubs = sub.split("-")
32	        assert len(subsubs) > 0
33	        if len(subsubs) == 1:
34	            indices.append(int(subsubs[0]))
35	        else:

148	                path = os.path.join(self.root, self.FILES[0])
149	                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
150	                    import academictorrents as at
151	                    atpath = at.get(self.AT_HASH, datastore=self.root)
152	                    assert atpath == path
153	
154	                print("Extracting {} to {}".format(path, datadir))
155	                os.makedirs(datadir, exist_ok=True)

206	                path = os.path.join(self.root, self.FILES[0])
207	                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
208	                    import academictorrents as at
209	                    atpath = at.get(self.AT_HASH, datastore=self.root)
210	                    assert atpath == path
211	
212	                print("Extracting {} to {}".format(path, datadir))
213	                os.makedirs(datadir, exist_ok=True)

270	    return preprocessor
271	
272	
273	def rgba_to_depth(x):
274	    assert x.dtype == np.uint8
275	    assert len(x.shape) == 3 and x.shape[2] == 4
276	    y = x.copy()
277	    y.dtype = np.float32

271	
272	
273	def rgba_to_depth(x):
274	    assert x.dtype == np.uint8
275	    assert len(x.shape) == 3 and x.shape[2] == 4
276	    y = x.copy()
277	    y.dtype = np.float32
278	    y = y.reshape(x.shape[:2])

417	    if factor is None or factor==1:
418	        return x
419	
420	    dtype = x.dtype
421	    assert dtype in [np.float32, np.float64]
422	    assert x.min() >= -1
423	    assert x.max() <= 1
424	

418	        return x
419	
420	    dtype = x.dtype
421	    assert dtype in [np.float32, np.float64]
422	    assert x.min() >= -1
423	    assert x.max() <= 1
424	
425	    keepmode = {"nearest": Image.NEAREST, "bilinear": Image.BILINEAR,

419	
420	    dtype = x.dtype
421	    assert dtype in [np.float32, np.float64]
422	    assert x.min() >= -1
423	    assert x.max() <= 1
424	
425	    keepmode = {"nearest": Image.NEAREST, "bilinear": Image.BILINEAR,
426	                "bicubic": Image.BICUBIC}[keepmode]

431	
432	    h, w, _ = x.shape
433	    nh = h//factor
434	    nw = w//factor
435	    assert nh > 0 and nw > 0, (nh, nw)
436	
437	    lr = lr.resize((nw,nh), Image.BICUBIC)
438	    if keepshapes:

60	        image = np.array(image).astype(np.uint8)
61	        if self.size is not None:
62	            image = self.image_rescaler(image=image)["image"]
63	        segmentation = Image.open(example["segmentation_path_"])
64	        assert segmentation.mode == "L", segmentation.mode
65	        segmentation = np.array(segmentation).astype(np.uint8)
66	        if self.shift_segmentation:
67	            # used to support segmentations containing unlabeled==255 label

12	    Build sinusoidal embeddings.
13	    This matches the implementation in tensor2tensor, but differs slightly
14	    from the description in Section 3.5 of "Attention Is All You Need".
15	    """
16	    assert len(timesteps.shape) == 1
17	
18	    half_dim = embedding_dim // 2
19	    emb = math.log(10000) / (half_dim - 1)

296	        #assert x.shape[2] == x.shape[3] == self.resolution
297	
298	        if self.use_timestep:
299	            # timestep embedding
300	            assert t is not None
301	            temb = get_timestep_embedding(t, self.ch)
302	            temb = self.temb.dense[0](temb)
303	            temb = nonlinearity(temb)

646	        #assert x.shape[2] == x.shape[3] == self.resolution
647	
648	        if self.use_timestep:
649	            # timestep embedding
650	            assert t is not None
651	            temb = get_timestep_embedding(t, self.ch)
652	            temb = self.temb.dense[0](temb)
653	            temb = nonlinearity(temb)

33	    return hashlib.md5(content).hexdigest()
34	
35	
36	def get_ckpt_path(name, root, check=False):
37	    assert name in URL_MAP
38	    path = os.path.join(root, CKPT_MAP[name])
39	    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
40	        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))

39	    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
40	        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
41	        download(URL_MAP[name], path)
42	        md5 = md5_hash(path)
43	        assert md5 == MD5_MAP[name], md5
44	    return path
45	
46	

36	                 disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
37	                 perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
38	                 disc_ndf=64, disc_loss="hinge"):
39	        super().__init__()
40	        assert disc_loss in ["hinge", "vanilla"]
41	        self.codebook_weight = codebook_weight
42	        self.pixel_weight = pixelloss_weight
43	        self.perceptual_loss = LPIPS().eval()

89	        # now the GAN part
90	        if optimizer_idx == 0:
91	            # generator update
92	            if cond is None:
93	                assert not self.disc_conditional
94	                logits_fake = self.discriminator(reconstructions.contiguous())
95	            else:
96	                assert self.disc_conditional

92	            if cond is None:
93	                assert not self.disc_conditional
94	                logits_fake = self.discriminator(reconstructions.contiguous())
95	            else:
96	                assert self.disc_conditional
97	                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1))
98	            g_loss = -torch.mean(logits_fake)
99	

99	
100	            try:
101	                d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
102	            except RuntimeError:
103	                assert not self.training
104	                d_weight = torch.tensor(0.0)
105	
106	            disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)

9	        return self
10	
11	    def encode(self, c):
12	        """fake vqmodel interface"""
13	        assert 0.0 <= c.min() and c.max() <= 1.0
14	        b,ch,h,w = c.shape
15	        assert ch == 1
16	

11	    def encode(self, c):
12	        """fake vqmodel interface"""
13	        assert 0.0 <= c.min() and c.max() <= 1.0
14	        b,ch,h,w = c.shape
15	        assert ch == 1
16	
17	        c = torch.nn.functional.interpolate(c, scale_factor=1/self.down_factor,
18	                                            mode="area")

9	
10	class ActNorm(nn.Module):
11	    def __init__(self, num_features, logdet=False, affine=True,
12	                 allow_reverse_init=False):
13	        assert affine
14	        super().__init__()
15	        self.logdet = logdet
16	        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))

145	            self.re_embed = n_embed
146	
147	    def remap_to_used(self, inds):
148	        ishape = inds.shape
149	        assert len(ishape)>1
150	        inds = inds.reshape(ishape[0],-1)
151	        used = self.used.to(inds)
152	        match = (inds[:,:,None]==used[None,None,...]).long()

159	        return new.reshape(ishape)
160	
161	    def unmap_to_all(self, inds):
162	        ishape = inds.shape
163	        assert len(ishape)>1
164	        inds = inds.reshape(ishape[0],-1)
165	        used = self.used.to(inds)
166	        if self.re_embed > self.used.shape[0]: # extra token

200	        return z_q, diff, ind
201	
202	    def get_codebook_entry(self, indices, shape):
203	        b, h, w, c = shape
204	        assert b*h*w == indices.shape[0]
205	        indices = rearrange(indices, '(b h w) -> b h w', b=b, h=h, w=w)
206	        if self.remap is not None:
207	            indices = self.unmap_to_all(indices)

245	        self.sane_index_shape = sane_index_shape
246	
247	    def remap_to_used(self, inds):
248	        ishape = inds.shape
249	        assert len(ishape)>1
250	        inds = inds.reshape(ishape[0],-1)
251	        used = self.used.to(inds)
252	        match = (inds[:,:,None]==used[None,None,...]).long()

259	        return new.reshape(ishape)
260	
261	    def unmap_to_all(self, inds):
262	        ishape = inds.shape
263	        assert len(ishape)>1
264	        inds = inds.reshape(ishape[0],-1)
265	        used = self.used.to(inds)
266	        if self.re_embed > self.used.shape[0]: # extra token

268	        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
269	        return back.reshape(ishape)
270	
271	    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
272	        assert temp is None or temp==1.0, "Only for interface compatible with Gumbel"
273	        assert rescale_logits==False, "Only for interface compatible with Gumbel"
274	        assert return_logits==False, "Only for interface compatible with Gumbel"
275	        # reshape z -> (batch, height, width, channel) and flatten

269	        return back.reshape(ishape)
270	
271	    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
272	        assert temp is None or temp==1.0, "Only for interface compatible with Gumbel"
273	        assert rescale_logits==False, "Only for interface compatible with Gumbel"
274	        assert return_logits==False, "Only for interface compatible with Gumbel"
275	        # reshape z -> (batch, height, width, channel) and flatten
276	        z = rearrange(z, 'b c h w -> b h w c').contiguous()

270	
271	    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
272	        assert temp is None or temp==1.0, "Only for interface compatible with Gumbel"
273	        assert rescale_logits==False, "Only for interface compatible with Gumbel"
274	        assert return_logits==False, "Only for interface compatible with Gumbel"
275	        # reshape z -> (batch, height, width, channel) and flatten
276	        z = rearrange(z, 'b c h w -> b h w c').contiguous()
277	        z_flattened = z.view(-1, self.e_dim)

381	            self.re_embed = n_embed
382	
383	    def remap_to_used(self, inds):
384	        ishape = inds.shape
385	        assert len(ishape)>1
386	        inds = inds.reshape(ishape[0],-1)
387	        used = self.used.to(inds)
388	        match = (inds[:,:,None]==used[None,None,...]).long()

395	        return new.reshape(ishape)
396	
397	    def unmap_to_all(self, inds):
398	        ishape = inds.shape
399	        assert len(ishape)>1
400	        inds = inds.reshape(ishape[0],-1)
401	        used = self.used.to(inds)
402	        if self.re_embed > self.used.shape[0]: # extra token

33	    return hashlib.md5(content).hexdigest()
34	
35	
36	def get_ckpt_path(name, root, check=False):
37	    assert name in URL_MAP
38	    path = os.path.join(root, CKPT_MAP[name])
39	    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
40	        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))

39	    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
40	        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
41	        download(URL_MAP[name], path)
42	        md5 = md5_hash(path)
43	        assert md5 == MD5_MAP[name], md5
44	    return path
45	
46	

67	    _OVERRIDE_FN[name] = fn
68	
69	
70	def update_override_act_fn(overrides):
71	    assert isinstance(overrides, dict)
72	    global _OVERRIDE_FN
73	    _OVERRIDE_FN.update(overrides)
74	

82	    _OVERRIDE_LAYER[name] = fn
83	
84	
85	def update_override_act_layer(overrides):
86	    assert isinstance(overrides, dict)
87	    global _OVERRIDE_LAYER
88	    _OVERRIDE_LAYER.update(overrides)
89	

141	    kwargs.setdefault('bias', False)
142	    padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs)
143	    if is_dynamic:
144	        if is_exportable():
145	            assert not is_scriptable()
146	            return Conv2dSameExport(in_chs, out_chs, kernel_size, **kwargs)
147	        else:
148	            return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs)

287	        return out
288	
289	
290	def select_conv2d(in_chs, out_chs, kernel_size, **kwargs):
291	    assert 'groups' not in kwargs  # only use 'depthwise' bool arg
292	    if isinstance(kernel_size, list):
293	        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently
294	        # We're going to use only lists for defining the MixedConv2d kernel groups,

289	
290	def select_conv2d(in_chs, out_chs, kernel_size, **kwargs):
291	    assert 'groups' not in kwargs  # only use 'depthwise' bool arg
292	    if isinstance(kernel_size, list):
293	        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently
294	        # We're going to use only lists for defining the MixedConv2d kernel groups,
295	        # ints, tuples, other iterables will continue to pass to normal conv and specify h, w.
296	        m = MixedConv2d(in_chs, out_chs, kernel_size, **kwargs)

58	    if not se_kwargs.pop('reduce_mid'):
59	        se_kwargs['reduced_base_chs'] = in_chs
60	    # act_layer override, if it remains None, the containing block's act_layer will be used
61	    if se_kwargs['act_layer'] is None:
62	        assert act_layer is not None
63	        se_kwargs['act_layer'] = act_layer
64	    return se_kwargs
65	

122	class ConvBnAct(nn.Module):
123	    def __init__(self, in_chs, out_chs, kernel_size,
124	                 stride=1, pad_type='', act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, norm_kwargs=None):
125	        super(ConvBnAct, self).__init__()
126	        assert stride in [1, 2]
127	        norm_kwargs = norm_kwargs or {}
128	        self.conv = select_conv2d(in_chs, out_chs, kernel_size, stride=stride, padding=pad_type)
129	        self.bn1 = norm_layer(out_chs, **norm_kwargs)

145	                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,
146	                 pw_kernel_size=1, pw_act=False, se_ratio=0., se_kwargs=None,
147	                 norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):
148	        super(DepthwiseSeparableConv, self).__init__()
149	        assert stride in [1, 2]
150	        norm_kwargs = norm_kwargs or {}
151	        self.has_residual = (stride == 1 and in_chs == out_chs) and not noskip
152	        self.drop_connect_rate = drop_connect_rate

396	        ba['norm_kwargs'] = self.norm_kwargs
397	        ba['pad_type'] = self.pad_type
398	        # block act fn overrides the model default
399	        ba['act_layer'] = ba['act_layer'] if ba['act_layer'] is not None else self.act_layer
400	        assert ba['act_layer'] is not None
401	        if bt == 'ir':
402	            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count
403	            ba['se_kwargs'] = self.se_kwargs

415	            block = EdgeResidual(**ba)
416	        elif bt == 'cn':
417	            block = ConvBnAct(**ba)
418	        else:
419	            assert False, 'Uknkown block type (%s) while building model.' % bt
420	        self.in_chs = ba['out_chs']  # update in_chs for arg of next block
421	        return block
422	

446	        self.block_idx = 0
447	        blocks = []
448	        # outer list of block_args defines the stacks ('stages' by some conventions)
449	        for stack_idx, stack in enumerate(block_args):
450	            assert isinstance(stack, list)
451	            stack = self._make_stack(stack)
452	            blocks.append(stack)
453	        return blocks

484	        A list of block args (dicts)
485	    Raises:
486	        ValueError: if the string def not properly specified (TODO)
487	    """
488	    assert isinstance(block_str, str)
489	    ops = block_str.split('_')
490	    block_type = ops[0]  # take the block type off the front
491	    ops = ops[1:]

573	            stride=int(options['s']),
574	            act_layer=act_layer,
575	        )
576	    else:
577	        assert False, 'Unknown block type (%s)' % block_type
578	
579	    return block_args, num_repeat
580	

619	
620	def decode_arch_def(arch_def, depth_multiplier=1.0, depth_trunc='ceil', experts_multiplier=1, fix_first_last=False):
621	    arch_args = []
622	    for stack_idx, block_strings in enumerate(arch_def):
623	        assert isinstance(block_strings, list)
624	        stack_args = []
625	        repeats = []
626	        for block_str in block_strings:

623	        assert isinstance(block_strings, list)
624	        stack_args = []
625	        repeats = []
626	        for block_str in block_strings:
627	            assert isinstance(block_str, str)
628	            ba, rep = _decode_block_str(block_str)
629	            if ba.get('num_experts', 0) > 0 and experts_multiplier > 1:
630	                ba['num_experts'] *= experts_multiplier

26	    graphs = []
27	    num_nodes = 0
28	    for node in graph.node:
29	        pn, gs = onnx.helper.printable_node(node, indent, subgraphs=True)
30	        assert isinstance(gs, list)
31	        content.append(pn)
32	        graphs.extend(gs)
33	        num_nodes += 1

44	        outdir_inc = outdir + '-' + str(count)
45	        while os.path.exists(outdir_inc):
46	            count = count + 1
47	            outdir_inc = outdir + '-' + str(count)
48	            assert count < 100
49	        outdir = outdir_inc
50	        os.makedirs(outdir)
51	    return outdir

86	    def decompress_keypoints(numbers: Optional[List[float]]) -> Optional[List[Optional[Keypoint]]]:
87	        if not numbers:
88	            return None
89	        
90	        assert len(numbers) % 3 == 0
91	
92	        def create_keypoint(x, y, c):
93	            if c < 1.0:

79	    def __call__(self, input_image, detect_resolution=512, safe=False, output_type="pil", scribble=False, upscale_method="INTER_CUBIC", **kwargs):
80	        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
81	        input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
82	
83	        assert input_image.ndim == 3
84	        H, W, C = input_image.shape
85	        with torch.no_grad():
86	            image_hed = torch.from_numpy(input_image).float().to(self.device)

108	
109	    Return an initialized network.
110	    """
111	    if len(gpu_ids) > 0:
112	        assert(torch.cuda.is_available())
113	        net.to(gpu_ids[0])
114	        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
115	    init_weights(net, init_type, init_gain=init_gain)

337	            use_dropout (bool)  -- if use dropout layers
338	            n_blocks (int)      -- the number of ResNet blocks
339	            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
340	        """
341	        assert(n_blocks >= 0)
342	        super(ResnetGenerator, self).__init__()
343	        if type(norm_layer) == functools.partial:
344	            use_bias = norm_layer.func == nn.InstanceNorm2d

121	        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
122	        detected_map, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
123	
124	        model = self.model_coarse if coarse else self.model
125	        assert detected_map.ndim == 3
126	        with torch.no_grad():
127	            image = torch.from_numpy(detected_map).float().to(self.device)
128	            image = image / 255.0

101	            refine_landmarks=True,
102	            min_detection_confidence=min_confidence,
103	    ) as facemesh:
104	        img_height, img_width, img_channels = img_rgb.shape
105	        assert(img_channels == 3)
106	
107	        results = facemesh.process(img_rgb).multi_face_landmarks
108	

126	                        #logger.info("Update config parameter {}: {} -> {}".format(param, config_param, arg_param))
127	                        setattr(config, param, arg_param)
128	
129	                # init a transformer encoder and append it to a list
130	                assert config.hidden_size % config.num_attention_heads == 0
131	                model = model_class(config=config) 
132	                #logger.info("Init model from scratch.")
133	                trans_encoder.append(model)

267	            tri_index = intersector.intersects_first(rays_o, rays_d)
268	
269	            tri_index = tri_index[index_ray]
270	
271	            assert len(index_ray) == len(tri_index)
272	            
273	            discriminator = (np.sum(mesh.face_normals[tri_index]* rays_d[index_ray], axis=-1)<= 0)
274	            points = points[discriminator] # ray intesects in interior faces, discard them

52	        resize_mode = "upper_bound"
53	        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
54	
55	    else:
56	        assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
57	
58	    transform = Compose(
59	        [

120	        )
121	
122	    else:
123	        print(f"model_type '{model_type}' not implemented, use: --model_type large")
124	        assert False
125	
126	    transform = Compose(
127	        [

156	    ]
157	
158	    def __init__(self, model_type, model_path):
159	        super().__init__()
160	        assert (model_type in self.MODEL_TYPES_ISL)
161	        model, _ = load_model(model_type, model_path)
162	        self.model = model
163	        self.model.train = disabled_train

39	                lines = pred_lines(img, self.model, [img.shape[0], img.shape[1]], thr_v, thr_d)
40	                for line in lines:
41	                    x_start, y_start, x_end, y_end = [int(val) for val in line]
42	                    cv2.line(img_output, (x_start, y_start), (x_end, y_end), [255, 255, 255], 1)
43	        except Exception as e:
44	            pass
45	
46	        detected_map = remove_pad(HWC3(img_output[:, :, 0]))
47	

124	class InvertedResidual(nn.Module):
125	    def __init__(self, inp, oup, stride, expand_ratio):
126	        super(InvertedResidual, self).__init__()
127	        self.stride = stride
128	        assert stride in [1, 2]
129	
130	        hidden_dim = int(round(inp * expand_ratio))
131	        self.use_res_connect = self.stride == 1 and inp == oup

123	class InvertedResidual(nn.Module):
124	    def __init__(self, inp, oup, stride, expand_ratio):
125	        super(InvertedResidual, self).__init__()
126	        self.stride = stride
127	        assert stride in [1, 2]
128	
129	        hidden_dim = int(round(inp * expand_ratio))
130	        self.use_res_connect = self.stride == 1 and inp == oup

22	    center: tpMap[1, 0, :, :]
23	    displacement: tpMap[1, 1:5, :, :]
24	    '''
25	    b, c, h, w = tpMap.shape
26	    assert  b==1, 'only support bsize==1'
27	    displacement = tpMap[:, 1:5, :, :][0]
28	    center = tpMap[:, 0, :, :]
29	    heat = torch.sigmoid(center)

551	        sorted_idx = np.argsort(score_array)[::-1]
552	        score_array = score_array[sorted_idx]
553	        squares = squares[sorted_idx]
554	
555	    except Exception as e:
556	        pass
557	
558	    '''return list
559	    merged_lines, squares, scores

67	    _OVERRIDE_FN[name] = fn
68	
69	
70	def update_override_act_fn(overrides):
71	    assert isinstance(overrides, dict)
72	    global _OVERRIDE_FN
73	    _OVERRIDE_FN.update(overrides)
74	

82	    _OVERRIDE_LAYER[name] = fn
83	
84	
85	def update_override_act_layer(overrides):
86	    assert isinstance(overrides, dict)
87	    global _OVERRIDE_LAYER
88	    _OVERRIDE_LAYER.update(overrides)
89	

141	    kwargs.setdefault('bias', False)
142	    padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs)
143	    if is_dynamic:
144	        if is_exportable():
145	            assert not is_scriptable()
146	            return Conv2dSameExport(in_chs, out_chs, kernel_size, **kwargs)
147	        else:
148	            return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs)

287	        return out
288	
289	
290	def select_conv2d(in_chs, out_chs, kernel_size, **kwargs):
291	    assert 'groups' not in kwargs  # only use 'depthwise' bool arg
292	    if isinstance(kernel_size, list):
293	        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently
294	        # We're going to use only lists for defining the MixedConv2d kernel groups,

289	
290	def select_conv2d(in_chs, out_chs, kernel_size, **kwargs):
291	    assert 'groups' not in kwargs  # only use 'depthwise' bool arg
292	    if isinstance(kernel_size, list):
293	        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently
294	        # We're going to use only lists for defining the MixedConv2d kernel groups,
295	        # ints, tuples, other iterables will continue to pass to normal conv and specify h, w.
296	        m = MixedConv2d(in_chs, out_chs, kernel_size, **kwargs)

58	    if not se_kwargs.pop('reduce_mid'):
59	        se_kwargs['reduced_base_chs'] = in_chs
60	    # act_layer override, if it remains None, the containing block's act_layer will be used
61	    if se_kwargs['act_layer'] is None:
62	        assert act_layer is not None
63	        se_kwargs['act_layer'] = act_layer
64	    return se_kwargs
65	

122	class ConvBnAct(nn.Module):
123	    def __init__(self, in_chs, out_chs, kernel_size,
124	                 stride=1, pad_type='', act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, norm_kwargs=None):
125	        super(ConvBnAct, self).__init__()
126	        assert stride in [1, 2]
127	        norm_kwargs = norm_kwargs or {}
128	        self.conv = select_conv2d(in_chs, out_chs, kernel_size, stride=stride, padding=pad_type)
129	        self.bn1 = norm_layer(out_chs, **norm_kwargs)

145	                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,
146	                 pw_kernel_size=1, pw_act=False, se_ratio=0., se_kwargs=None,
147	                 norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):
148	        super(DepthwiseSeparableConv, self).__init__()
149	        assert stride in [1, 2]
150	        norm_kwargs = norm_kwargs or {}
151	        self.has_residual = (stride == 1 and in_chs == out_chs) and not noskip
152	        self.drop_connect_rate = drop_connect_rate

396	        ba['norm_kwargs'] = self.norm_kwargs
397	        ba['pad_type'] = self.pad_type
398	        # block act fn overrides the model default
399	        ba['act_layer'] = ba['act_layer'] if ba['act_layer'] is not None else self.act_layer
400	        assert ba['act_layer'] is not None
401	        if bt == 'ir':
402	            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count
403	            ba['se_kwargs'] = self.se_kwargs

415	            block = EdgeResidual(**ba)
416	        elif bt == 'cn':
417	            block = ConvBnAct(**ba)
418	        else:
419	            assert False, 'Uknkown block type (%s) while building model.' % bt
420	        self.in_chs = ba['out_chs']  # update in_chs for arg of next block
421	        return block
422	

446	        self.block_idx = 0
447	        blocks = []
448	        # outer list of block_args defines the stacks ('stages' by some conventions)
449	        for stack_idx, stack in enumerate(block_args):
450	            assert isinstance(stack, list)
451	            stack = self._make_stack(stack)
452	            blocks.append(stack)
453	        return blocks

484	        A list of block args (dicts)
485	    Raises:
486	        ValueError: if the string def not properly specified (TODO)
487	    """
488	    assert isinstance(block_str, str)
489	    ops = block_str.split('_')
490	    block_type = ops[0]  # take the block type off the front
491	    ops = ops[1:]

573	            stride=int(options['s']),
574	            act_layer=act_layer,
575	        )
576	    else:
577	        assert False, 'Unknown block type (%s)' % block_type
578	
579	    return block_args, num_repeat
580	

619	
620	def decode_arch_def(arch_def, depth_multiplier=1.0, depth_trunc='ceil', experts_multiplier=1, fix_first_last=False):
621	    arch_args = []
622	    for stack_idx, block_strings in enumerate(arch_def):
623	        assert isinstance(block_strings, list)
624	        stack_args = []
625	        repeats = []
626	        for block_str in block_strings:

623	        assert isinstance(block_strings, list)
624	        stack_args = []
625	        repeats = []
626	        for block_str in block_strings:
627	            assert isinstance(block_str, str)
628	            ba, rep = _decode_block_str(block_str)
629	            if ba.get('num_experts', 0) > 0 and experts_multiplier > 1:
630	                ba['num_experts'] *= experts_multiplier

26	    graphs = []
27	    num_nodes = 0
28	    for node in graph.node:
29	        pn, gs = onnx.helper.printable_node(node, indent, subgraphs=True)
30	        assert isinstance(gs, list)
31	        content.append(pn)
32	        graphs.extend(gs)
33	        num_nodes += 1

44	        outdir_inc = outdir + '-' + str(count)
45	        while os.path.exists(outdir_inc):
46	            count = count + 1
47	            outdir_inc = outdir + '-' + str(count)
48	            assert count < 100
49	        outdir = outdir_inc
50	        os.makedirs(outdir)
51	    return outdir

295	        },
296	    }
297	
298	def createConvFunc(op_type):
299	    assert op_type in ['cv', 'cd', 'ad', 'rd'], 'unknown op type: %s' % str(op_type)
300	    if op_type == 'cv':
301	        return F.conv2d
302	

301	        return F.conv2d
302	
303	    if op_type == 'cd':
304	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
305	            assert dilation in [1, 2], 'dilation for cd_conv should be in 1 or 2'
306	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for cd_conv should be 3x3'
307	            assert padding == dilation, 'padding for cd_conv set wrong'
308	

302	
303	    if op_type == 'cd':
304	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
305	            assert dilation in [1, 2], 'dilation for cd_conv should be in 1 or 2'
306	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for cd_conv should be 3x3'
307	            assert padding == dilation, 'padding for cd_conv set wrong'
308	
309	            weights_c = weights.sum(dim=[2, 3], keepdim=True)

303	    if op_type == 'cd':
304	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
305	            assert dilation in [1, 2], 'dilation for cd_conv should be in 1 or 2'
306	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for cd_conv should be 3x3'
307	            assert padding == dilation, 'padding for cd_conv set wrong'
308	
309	            weights_c = weights.sum(dim=[2, 3], keepdim=True)
310	            yc = F.conv2d(x, weights_c, stride=stride, padding=0, groups=groups)

312	            return y - yc
313	        return func
314	    elif op_type == 'ad':
315	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
316	            assert dilation in [1, 2], 'dilation for ad_conv should be in 1 or 2'
317	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for ad_conv should be 3x3'
318	            assert padding == dilation, 'padding for ad_conv set wrong'
319	

313	        return func
314	    elif op_type == 'ad':
315	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
316	            assert dilation in [1, 2], 'dilation for ad_conv should be in 1 or 2'
317	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for ad_conv should be 3x3'
318	            assert padding == dilation, 'padding for ad_conv set wrong'
319	
320	            shape = weights.shape

314	    elif op_type == 'ad':
315	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
316	            assert dilation in [1, 2], 'dilation for ad_conv should be in 1 or 2'
317	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for ad_conv should be 3x3'
318	            assert padding == dilation, 'padding for ad_conv set wrong'
319	
320	            shape = weights.shape
321	            weights = weights.view(shape[0], shape[1], -1)

324	            return y
325	        return func
326	    elif op_type == 'rd':
327	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
328	            assert dilation in [1, 2], 'dilation for rd_conv should be in 1 or 2'
329	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for rd_conv should be 3x3'
330	            padding = 2 * dilation
331	

325	        return func
326	    elif op_type == 'rd':
327	        def func(x, weights, bias=None, stride=1, padding=0, dilation=1, groups=1):
328	            assert dilation in [1, 2], 'dilation for rd_conv should be in 1 or 2'
329	            assert weights.size(2) == 3 and weights.size(3) == 3, 'kernel size for rd_conv should be 3x3'
330	            padding = 2 * dilation
331	
332	            shape = weights.shape

497	    def __init__(self, inplane, pdcs, dil=None, sa=False, convert=False):
498	        super(PiDiNet, self).__init__()
499	        self.sa = sa
500	        if dil is not None:
501	            assert isinstance(dil, int), 'dil should be an int'
502	        self.dil = dil
503	
504	        self.fuseplanes = []

647	        return outputs
648	
649	def config_model(model):
650	    model_options = list(nets.keys())
651	    assert model in model_options, \
652	        'unrecognized model, please choose from %s' % str(model_options)
653	
654	    # print(str(nets[model]))
655	

30	class Recolorizer:
31	    def __call__(self, input_image=None, mode="luminance", gamma_correction=1.0, detect_resolution=512, output_type=None, upscale_method="INTER_CUBIC", **kwargs):
32	        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
33	        input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
34	        assert mode in recolor_methods.keys()
35	        detected_map = recolor_methods[mode](input_image, gamma_correction)
36	        detected_map = HWC3(remove_pad(detected_map))
37	        if output_type == "pil":

94	            For large resolutions, 'binary_mask' may consume large amounts of
95	            memory.
96	        """
97	
98	        assert (points_per_side is None) != (
99	            point_grids is None
100	        ), "Exactly one of points_per_side or point_grid must be provided."
101	        if points_per_side is not None:
102	            self.point_grids = build_all_layer_point_grids(
103	                points_per_side,

108	            self.point_grids = point_grids
109	        else:
110	            raise ValueError("Can't have both points_per_side and point_grid be None.")
111	
112	        assert output_mode in [
113	            "binary_mask",
114	            "uncompressed_rle",
115	            "coco_rle",
116	        ], f"Unknown output_mode {output_mode}."
117	        if output_mode == "coco_rle":
118	            from custom_pycocotools import mask as mask_utils  # type: ignore # noqa: F401
119	

213	        self.proj = nn.Linear(dim, dim)
214	
215	        self.use_rel_pos = use_rel_pos
216	        if self.use_rel_pos:
217	            assert (
218	                input_size is not None
219	            ), "Input size must be provided if using relative positional encoding."
220	            # initialize relative positional embeddings
221	            self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
222	            self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))

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	

636	    def fn_wrapper(pretrained=False, **kwargs):
637	        model = fn()
638	        if pretrained:
639	            model_name = fn.__name__
640	            assert model_name in _provided_checkpoints, \
641	                f'Sorry that the checkpoint `{model_name}` is not provided yet.'
642	            url = _checkpoint_url_format.format(
643	                _provided_checkpoints[model_name])
644	            checkpoint = torch.hub.load_state_dict_from_url(

197	        super().__init__()
198	        self.embedding_dim = embedding_dim
199	        self.internal_dim = embedding_dim // downsample_rate
200	        self.num_heads = num_heads
201	        assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."
202	
203	        self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
204	        self.k_proj = nn.Linear(embedding_dim, self.internal_dim)

44	          image (np.ndarray): The image for calculating masks. Expects an
45	            image in HWC uint8 format, with pixel values in [0, 255].
46	          image_format (str): The color format of the image, in ['RGB', 'BGR'].
47	        """
48	        assert image_format in [
49	            "RGB",
50	            "BGR",
51	        ], f"image_format must be in ['RGB', 'BGR'], is {image_format}."
52	        if image_format != self.model.image_format:
53	            image = image[..., ::-1]
54	

75	            1x3xHxW, which has been transformed with ResizeLongestSide.
76	          original_image_size (tuple(int, int)): The size of the image
77	            before transformation, in (H, W) format.
78	        """
79	        assert (
80	            len(transformed_image.shape) == 4
81	            and transformed_image.shape[1] == 3
82	            and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size
83	        ), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}."
84	        self.reset_image()
85	
86	        self.original_size = original_image_size

135	
136	        # Transform input prompts
137	        coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None
138	        if point_coords is not None:
139	            assert (
140	                point_labels is not None
141	            ), "point_labels must be supplied if point_coords is supplied."
142	            point_coords = self.transform.apply_coords(point_coords, self.original_size)
143	            coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device)
144	            labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)

251	        if not self.is_image_set:
252	            raise RuntimeError(
253	                "An image must be set with .set_image(...) to generate an embedding."
254	            )
255	        assert self.features is not None, "Features must exist if an image has been set."
256	        return self.features
257	
258	    @property

20	    """
21	
22	    def __init__(self, **kwargs) -> None:
23	        for v in kwargs.values():
24	            assert isinstance(
25	                v, (list, np.ndarray, torch.Tensor)
26	            ), "MaskData only supports list, numpy arrays, and torch tensors."
27	        self._stats = dict(**kwargs)
28	
29	    def __setitem__(self, key: str, item: Any) -> None:

26	            ), "MaskData only supports list, numpy arrays, and torch tensors."
27	        self._stats = dict(**kwargs)
28	
29	    def __setitem__(self, key: str, item: Any) -> None:
30	        assert isinstance(
31	            item, (list, np.ndarray, torch.Tensor)
32	        ), "MaskData only supports list, numpy arrays, and torch tensors."
33	        self._stats[key] = item
34	
35	    def __delitem__(self, key: str) -> None:

95	    return box_xywh
96	
97	
98	def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
99	    assert len(args) > 0 and all(
100	        len(a) == len(args[0]) for a in args
101	    ), "Batched iteration must have inputs of all the same size."
102	    n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)
103	    for b in range(n_batches):
104	        yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]

272	    mask and an indicator of if the mask has been modified.
273	    """
274	    import cv2  # type: ignore
275	
276	    assert mode in ["holes", "islands"]
277	    correct_holes = mode == "holes"
278	    working_mask = (correct_holes ^ mask).astype(np.uint8)
279	    n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8)

54	    def __call__(self, img):
55	        eps = 1e-5
56	        X = img.astype(np.float32)
57	        r, g, b = X[:, :, 0], X[:, :, 1], X[:, :, 2]
58	        kr, kg, kb = [random.random() + eps for _ in range(3)]
59	        ks = kr + kg + kb
60	        kr /= ks
61	        kg /= ks

119	        self.up_factor = 2
120	        self.constant_features = 16
121	
122	        layers = self.make_deconv_layers(in_features, up_scale)
123	        assert layers is not None, layers
124	        self.features = nn.Sequential(*layers)
125	
126	    def make_deconv_layers(self, in_features, up_scale):

237	
238	    def crop_bdcn(data1, h, w, crop_h, crop_w):
239	        # Based on BDCN Implementation @ https://github.com/pkuCactus/BDCN
240	        _, _, h1, w1 = data1.size()
241	        assert (h <= h1 and w <= w1)
242	        data = data1[:, :, crop_h:crop_h + h, crop_w:crop_w + w]
243	        return data
244	

243	        return data
244	
245	
246	    def forward(self, x, single_test=False):
247	        assert x.ndim == 4, x.shape
248	         # supose the image size is 352x352
249	
250	        # Block 1

39	    return Processor(request.param)
40	
41	
42	def test_processor_init(processor):
43	    assert isinstance(processor.processor, MODELS[processor.processor_id]['class'])
44	    assert isinstance(processor.params, dict)
45	
46	

40	
41	
42	def test_processor_init(processor):
43	    assert isinstance(processor.processor, MODELS[processor.processor_id]['class'])
44	    assert isinstance(processor.params, dict)
45	
46	
47	def test_processor_call(processor):

62	    W = int(np.round(W / 64.0)) * 64
63	
64	    # Test processing
65	    processed_image = processor(image)
66	    assert isinstance(processed_image, Image.Image)
67	    assert processed_image.size == (W, H)
68	
69	

63	
64	    # Test processing
65	    processed_image = processor(image)
66	    assert isinstance(processed_image, Image.Image)
67	    assert processed_image.size == (W, H)
68	
69	
70	def test_processor_call_bytes(processor):

73	        image_bytes = f.read()
74	
75	    # Test processing
76	    processed_image_bytes = processor(image_bytes, to_pil=False)
77	    assert isinstance(processed_image_bytes, bytes)
78	    assert len(processed_image_bytes) > 0

74	
75	    # Test processing
76	    processed_image_bytes = processor(image_bytes, to_pil=False)
77	    assert isinstance(processed_image_bytes, bytes)
78	    assert len(processed_image_bytes) > 0

184	def get_external_models():
185	    mmcv_home = _get_mmcv_home()
186	    default_json_path = osp.join(mmcv.__path__[0], 'model_zoo/open_mmlab.json')
187	    default_urls = load_file(default_json_path)
188	    assert isinstance(default_urls, dict)
189	    external_json_path = osp.join(mmcv_home, 'open_mmlab.json')
190	    if osp.exists(external_json_path):
191	        external_urls = load_file(external_json_path)

188	    assert isinstance(default_urls, dict)
189	    external_json_path = osp.join(mmcv_home, 'open_mmlab.json')
190	    if osp.exists(external_json_path):
191	        external_urls = load_file(external_json_path)
192	        assert isinstance(external_urls, dict)
193	        default_urls.update(external_urls)
194	
195	    return default_urls

205	def get_deprecated_model_names():
206	    deprecate_json_path = osp.join(mmcv.__path__[0],
207	                                   'model_zoo/deprecated.json')
208	    deprecate_urls = load_file(deprecate_json_path)
209	    assert isinstance(deprecate_urls, dict)
210	
211	    return deprecate_urls
212	

43	    nearest_size = [int(np.ceil(image1.size(-2) / padding_factor)) * padding_factor,
44	                    int(np.ceil(image1.size(-1) / padding_factor)) * padding_factor]
45	    # resize to nearest size or specified size
46	    inference_size = nearest_size if fixed_inference_size is None else fixed_inference_size
47	    assert isinstance(inference_size, list) or isinstance(inference_size, tuple)
48	    ori_size = image1.shape[-2:]
49	
50	    # resize before inference

82	    vis_image = flow_to_image(flow)
83	
84	    # also predict backward flow
85	    if pred_bidir_flow:
86	        assert flow_pr.size(0) == 2  # [2, H, W, 2]
87	        flow_bwd = flow_pr[1].permute(1, 2, 0).cpu().numpy()  # [H, W, 2]
88	        vis_image = flow_to_image(flow_bwd)
89	        flow = flow_bwd

158	        for key in list(MODEL_CONFIGS.keys())[::-1]:
159	            if key in filename:
160	                config_args = MODEL_CONFIGS[key]
161	                break
162	        assert config_args, f"Couldn't find hardcoded Unimatch config for {filename}"
163	        
164	        model = UniMatch(feature_channels=config_args.feature_channels,
165	                        num_scales=config_args.num_scales,

179	        self.device = device
180	        return self
181	    
182	    def __call__(self, image1, image2, detect_resolution=512, output_type="pil", upscale_method="INTER_CUBIC", pred_bwd_flow=False, pred_bidir_flow=False, **kwargs):
183	        assert image1.shape == image2.shape, f"[Unimatch] image1 and image2 must have the same size, got {image1.shape} and {image2.shape}"
184	
185	        image1, output_type = common_input_validate(image1, output_type, **kwargs)
186	        #image1, remove_pad = resize_image_with_pad(image1, detect_resolution, upscale_method)

6	
7	
8	def single_head_full_attention(q, k, v):
9	    # q, k, v: [B, L, C]
10	    assert q.dim() == k.dim() == v.dim() == 3
11	
12	    scores = torch.matmul(q, k.permute(0, 2, 1)) / (q.size(2) ** .5)  # [B, L, L]
13	    attn = torch.softmax(scores, dim=2)  # [B, L, L]

21	                                  w=None,
22	                                  ):
23	    # q, k, v: [B, L, C]
24	
25	    assert h is not None and w is not None
26	    assert q.size(1) == h * w
27	
28	    b, _, c = q.size()

22	                                  ):
23	    # q, k, v: [B, L, C]
24	
25	    assert h is not None and w is not None
26	    assert q.size(1) == h * w
27	
28	    b, _, c = q.size()
29	

50	                                       attn_mask=None,
51	                                       ):
52	    # ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
53	    # q, k, v: [B, L, C]
54	    assert q.dim() == k.dim() == v.dim() == 3
55	
56	    assert h is not None and w is not None
57	    assert q.size(1) == h * w

52	    # ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
53	    # q, k, v: [B, L, C]
54	    assert q.dim() == k.dim() == v.dim() == 3
55	
56	    assert h is not None and w is not None
57	    assert q.size(1) == h * w
58	
59	    b, _, c = q.size()

53	    # q, k, v: [B, L, C]
54	    assert q.dim() == k.dim() == v.dim() == 3
55	
56	    assert h is not None and w is not None
57	    assert q.size(1) == h * w
58	
59	    b, _, c = q.size()
60	

69	
70	    scale_factor = c ** 0.5
71	
72	    if with_shift:
73	        assert attn_mask is not None  # compute once
74	        shift_size_h = window_size_h // 2
75	        shift_size_w = window_size_w // 2
76	

113	                                          attn_mask=None,
114	                                          ):
115	    # q, k, v: [B, L, C]
116	
117	    assert h is not None and w is not None
118	    assert q.size(1) == h * w
119	
120	    b, _, c = q.size()

114	                                          ):
115	    # q, k, v: [B, L, C]
116	
117	    assert h is not None and w is not None
118	    assert q.size(1) == h * w
119	
120	    b, _, c = q.size()
121	

129	
130	    scale_factor = c ** 0.5
131	
132	    if with_shift:
133	        assert attn_mask is not None  # compute once
134	        shift_size_w = window_size_w // 2
135	
136	        q = torch.roll(q, shifts=-shift_size_w, dims=1)

216	
217	    def forward_local_window_attn(self, feature0, flow,
218	                                  local_window_radius=1,
219	                                  ):
220	        assert flow.size(1) == 2 or flow.size(1) == 1  # flow or disparity or depth
221	        assert local_window_radius > 0
222	
223	        b, c, h, w = feature0.size()

217	    def forward_local_window_attn(self, feature0, flow,
218	                                  local_window_radius=1,
219	                                  ):
220	        assert flow.size(1) == 2 or flow.size(1) == 1  # flow or disparity or depth
221	        assert local_window_radius > 0
222	
223	        b, c, h, w = feature0.size()
224	

21	    return grid
22	
23	
24	def generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None):
25	    assert device is not None
26	
27	    x, y = torch.meshgrid([torch.linspace(w_min, w_max, len_w, device=device),
28	                           torch.linspace(h_min, h_max, len_h, device=device)],

63	
64	
65	def flow_warp(feature, flow, mask=False, padding_mode='zeros'):
66	    b, c, h, w = feature.size()
67	    assert flow.size(1) == 2
68	
69	    grid = coords_grid(b, h, w).to(flow.device) + flow  # [B, 2, H, W]
70	

77	                                       beta=0.5
78	                                       ):
79	    # fwd_flow, bwd_flow: [B, 2, H, W]
80	    # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
81	    assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
82	    assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
83	    flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1)  # [B, H, W]
84	

78	                                       ):
79	    # fwd_flow, bwd_flow: [B, 2, H, W]
80	    # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
81	    assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
82	    assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
83	    flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1)  # [B, H, W]
84	
85	    warped_bwd_flow = flow_warp(bwd_flow, fwd_flow)  # [B, 2, H, W]

106	    if not isinstance(flow, float):
107	        sample_coords = sample_coords + flow.view(
108	            b, 2, -1).permute(0, 2, 1).unsqueeze(-2)  # [B, H*W, (2R+1)^2, 2]
109	    else:
110	        assert flow == 0.
111	
112	    # normalize coordinates to [-1, 1]
113	    sample_coords_norm = normalize_coords(sample_coords, h, w)  # [-1, 1]

207	                              depth_from_argmax=False,
208	                              pred_bidir_depth=False,
209	                              ):
210	    b, c, h, w = feature0.size()
211	    assert depth_candidates.dim() == 4  # [B, D, H, W]
212	    scale_factor = c ** 0.5
213	
214	    if pred_bidir_depth:

245	    pose: [B, 4, 4]
246	    depth: [B, D, H, W]
247	    """
248	
249	    assert intrinsics.size(1) == intrinsics.size(2) == 3
250	    assert pose.size(1) == pose.size(2) == 4
251	    assert depth.dim() == 4
252	

246	    depth: [B, D, H, W]
247	    """
248	
249	    assert intrinsics.size(1) == intrinsics.size(2) == 3
250	    assert pose.size(1) == pose.size(2) == 4
251	    assert depth.dim() == 4
252	
253	    b, d, h, w = depth.size()

247	    """
248	
249	    assert intrinsics.size(1) == intrinsics.size(2) == 3
250	    assert pose.size(1) == pose.size(2) == 4
251	    assert depth.dim() == 4
252	
253	    b, d, h, w = depth.size()
254	    c = feature1.size(1)

115	                        # full 2d attn
116	                        message = single_head_full_attention(query, key, value)  # [N, L, C]
117	                else:
118	                    if attn_num_splits > 1:
119	                        assert shifted_window_attn_mask_1d is not None
120	                        # cross attn 1d shift
121	                        message = single_head_split_window_attention_1d(query, key, value,
122	                                                                        num_splits=attn_num_splits,

229	                **kwargs,
230	                ):
231	
232	        b, c, h, w = feature0.shape
233	        assert self.d_model == c
234	
235	        feature0 = feature0.flatten(-2).permute(0, 2, 1)  # [B, H*W, C]
236	        feature1 = feature1.flatten(-2).permute(0, 2, 1)  # [B, H*W, C]

46	        self.test_branch_idx = test_branch_idx
47	        self.norm = norm
48	        self.activation = activation
49	
50	        assert len({self.num_branch, len(self.paddings), len(self.strides)}) == 1
51	
52	        self.weight = nn.Parameter(
53	            torch.Tensor(out_channels, in_channels // groups, *self.kernel_size)

62	            nn.init.constant_(self.bias, 0)
63	
64	    def forward(self, inputs):
65	        num_branch = self.num_branch if self.training or self.test_branch_idx == -1 else 1
66	        assert len(inputs) == num_branch
67	
68	        if self.training or self.test_branch_idx == -1:
69	            outputs = [

110	                **kwargs,
111	                ):
112	
113	        if pred_bidir_flow:
114	            assert task == 'flow'
115	
116	        if task == 'depth':
117	            assert self.num_scales == 1  # multi-scale depth model is not supported yet

113	        if pred_bidir_flow:
114	            assert task == 'flow'
115	
116	        if task == 'depth':
117	            assert self.num_scales == 1  # multi-scale depth model is not supported yet
118	
119	        results_dict = {}
120	        flow_preds = []

128	
129	        flow = None
130	
131	        if task != 'depth':
132	            assert len(attn_splits_list) == len(corr_radius_list) == len(prop_radius_list) == self.num_scales
133	        else:
134	            assert len(attn_splits_list) == len(prop_radius_list) == self.num_scales == 1
135	

130	
131	        if task != 'depth':
132	            assert len(attn_splits_list) == len(corr_radius_list) == len(prop_radius_list) == self.num_scales
133	        else:
134	            assert len(attn_splits_list) == len(prop_radius_list) == self.num_scales == 1
135	
136	        for scale_idx in range(self.num_scales):
137	            feature0, feature1 = feature0_list[scale_idx], feature1_list[scale_idx]

149	                intrinsics_curr = intrinsics.clone()
150	                intrinsics_curr[:, :2] = intrinsics_curr[:, :2] / upsample_factor
151	
152	            if scale_idx > 0:
153	                assert task != 'depth'  # not supported for multi-scale depth model
154	                flow = F.interpolate(flow, scale_factor=2, mode='bilinear', align_corners=True) * 2
155	
156	            if flow is not None:

153	                assert task != 'depth'  # not supported for multi-scale depth model
154	                flow = F.interpolate(flow, scale_factor=2, mode='bilinear', align_corners=True) * 2
155	
156	            if flow is not None:
157	                assert task != 'depth'
158	                flow = flow.detach()
159	
160	                if task == 'stereo':

268	                                                     upsample_factor=upsample_factor,
269	                                                     is_depth=task == 'depth')
270	                        flow_preds.append(flow_up)
271	
272	                    assert num_reg_refine > 0
273	                    for refine_iter_idx in range(num_reg_refine):
274	                        flow = flow.detach()
275	

3	from .position import PositionEmbeddingSine
4	
5	
6	def generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None):
7	    assert device is not None
8	
9	    x, y = torch.meshgrid([torch.linspace(w_min, w_max, len_w, device=device),
10	                           torch.linspace(h_min, h_max, len_h, device=device)],

36	                  channel_last=False,
37	                  ):
38	    if channel_last:  # [B, H, W, C]
39	        b, h, w, c = feature.size()
40	        assert h % num_splits == 0 and w % num_splits == 0
41	
42	        b_new = b * num_splits * num_splits
43	        h_new = h // num_splits

46	        feature = feature.view(b, num_splits, h // num_splits, num_splits, w // num_splits, c
47	                               ).permute(0, 1, 3, 2, 4, 5).reshape(b_new, h_new, w_new, c)  # [B*K*K, H/K, W/K, C]
48	    else:  # [B, C, H, W]
49	        b, c, h, w = feature.size()
50	        assert h % num_splits == 0 and w % num_splits == 0
51	
52	        b_new = b * num_splits * num_splits
53	        h_new = h // num_splits

156	                     num_splits=2,
157	                     ):
158	    # feature: [B, W, C]
159	    b, w, c = feature.size()
160	    assert w % num_splits == 0
161	
162	    b_new = b * num_splits
163	    w_new = w // num_splits

1	# Copyright (c) OpenMMLab. All rights reserved.
2	# https://github.com/open-mmlab/mmcv/blob/7540cf73ac7e5d1e14d0ffbd9b6759e83929ecfc/mmcv/runner/dist_utils.py
3	
4	import os
5	import subprocess
6	
7	import torch
8	import torch.multiprocessing as mp

201	
202	
203	def extract_video(video_name):
204	    cap = cv2.VideoCapture(video_name)
205	    assert cap.isOpened(), f'Failed to load video file {video_name}'
206	    # get video info
207	    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
208	            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))

117	    :param clip_flow: float, maximum clipping value for flow
118	    :return:
119	    '''
120	
121	    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
122	    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
123	
124	    if clip_flow is not None:

118	    :return:
119	    '''
120	
121	    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
122	    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
123	
124	    if clip_flow is not None:
125	        flow_uv = np.clip(flow_uv, 0, clip_flow)

76	    """
77	    nBands = 2
78	
79	    if v is None:
80	        assert (uv.ndim == 3)
81	        assert (uv.shape[2] == 2)
82	        u = uv[:, :, 0]
83	        v = uv[:, :, 1]

77	    nBands = 2
78	
79	    if v is None:
80	        assert (uv.ndim == 3)
81	        assert (uv.shape[2] == 2)
82	        u = uv[:, :, 0]
83	        v = uv[:, :, 1]
84	    else:

83	        v = uv[:, :, 1]
84	    else:
85	        u = uv
86	
87	    assert (u.shape == v.shape)
88	    height, width = u.shape
89	    f = open(filename, 'wb')
90	    # write the header

143	    # In G, flow along x-axis normalized by image width and quantized to [0;2^16 – 1]
144	    # B = 0 for invalid flow (e.g., sky pixels)
145	    bgr = cv2.imread(filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
146	    h, w, _c = bgr.shape
147	    assert bgr.dtype == np.uint16 and _c == 3
148	    # b == invalid flow flag == 0 for sky or other invalid flow
149	    invalid = bgr[:, :, 0] == 0
150	    # g,r == flow_y,x normalized by height,width and scaled to [0;2**16 – 1]

36	    xgrid, ygrid = coords.split([1, 1], dim=-1)
37	
38	    # To handle H or W equals to 1 by explicitly defining height and width
39	    if H == 1:
40	        assert ygrid.abs().max() < 1e-8
41	        H = 10
42	    if W == 1:
43	        assert xgrid.abs().max() < 1e-8

39	    if H == 1:
40	        assert ygrid.abs().max() < 1e-8
41	        H = 10
42	    if W == 1:
43	        assert xgrid.abs().max() < 1e-8
44	        W = 10
45	
46	    xgrid = 2 * xgrid / (W - 1) - 1

77	
78	
79	def compute_out_of_boundary_mask(flow, downsample_factor=None):
80	    # flow: [B, 2, H, W]
81	    assert flow.dim() == 4 and flow.size(1) == 2
82	    b, _, h, w = flow.shape
83	    init_coords = coords_grid(b, h, w).to(flow.device)
84	    corres = init_coords + flow  # [B, 2, H, W]

83	    init_coords = coords_grid(b, h, w).to(flow.device)
84	    corres = init_coords + flow  # [B, 2, H, W]
85	
86	    if downsample_factor is not None:
87	        assert w % downsample_factor == 0 and h % downsample_factor == 0
88	        # the actual max disp can predict is in the downsampled feature resolution, then upsample
89	        max_w = (w // downsample_factor - 1) * downsample_factor
90	        max_h = (h // downsample_factor - 1) * downsample_factor

107	    """Normalize coordinates of image scale to [-1, 1]
108	    Args:
109	        grid: [B, 2, H, W]
110	    """
111	    assert grid.size(1) == 2
112	    h, w = grid.size()[2:]
113	    grid[:, 0, :, :] = 2 * (grid[:, 0, :, :].clone() / (w - 1)) - 1  # x: [-1, 1]
114	    grid[:, 1, :, :] = 2 * (grid[:, 1, :, :].clone() / (h - 1)) - 1  # y: [-1, 1]

117	
118	
119	def flow_warp(feature, flow, mask=False, padding_mode='zeros'):
120	    b, c, h, w = feature.size()
121	    assert flow.size(1) == 2
122	
123	    grid = coords_grid(b, h, w).to(flow.device) + flow  # [B, 2, H, W]
124	

130	    return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
131	
132	
133	def bilinear_upflow(flow, scale_factor=8):
134	    assert flow.size(1) == 2
135	    flow = F.interpolate(flow, scale_factor=scale_factor,
136	                         mode='bilinear', align_corners=True) * scale_factor
137	

90	
91	
92	def viz_depth_tensor(disp, return_numpy=False, colormap='plasma'):
93	    # visualize inverse depth
94	    assert isinstance(disp, torch.Tensor)
95	
96	    disp = disp.numpy()
97	    vmax = np.percentile(disp, 95)

51	
52	here = Path(__file__).parent.resolve()
53	
54	def HWC3(x):
55	    assert x.dtype == np.uint8
56	    if x.ndim == 2:
57	        x = x[:, :, None]
58	    assert x.ndim == 3

54	def HWC3(x):
55	    assert x.dtype == np.uint8
56	    if x.ndim == 2:
57	        x = x[:, :, None]
58	    assert x.ndim == 3
59	    H, W, C = x.shape
60	    assert C == 1 or C == 3 or C == 4
61	    if C == 3:

56	    if x.ndim == 2:
57	        x = x[:, :, None]
58	    assert x.ndim == 3
59	    H, W, C = x.shape
60	    assert C == 1 or C == 3 or C == 4
61	    if C == 3:
62	        return x
63	    if C == 1:

113	    return y
114	
115	
116	def img2mask(img, H, W, low=10, high=90):
117	    assert img.ndim == 3 or img.ndim == 2
118	    assert img.dtype == np.uint8
119	
120	    if img.ndim == 3:

114	
115	
116	def img2mask(img, H, W, low=10, high=90):
117	    assert img.ndim == 3 or img.ndim == 2
118	    assert img.dtype == np.uint8
119	
120	    if img.ndim == 3:
121	        y = img[:, :, random.randrange(0, img.shape[2])]

117	    assert img.ndim == 3 or img.ndim == 2
118	    assert img.dtype == np.uint8
119	
120	    if img.ndim == 3:
121	        y = img[:, :, random.randrange(0, img.shape[2])]
122	    else:
123	        y = img
124	

123	        y = img
124	
125	    y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC)
126	
127	    if random.uniform(0, 1) < 0.5:
128	        y = 255 - y
129	
130	    return y < np.percentile(y, random.randrange(low, high))